feat(api): allow listening with unix socket

This commit is contained in:
Guanran Wang 2024-11-02 12:07:02 +08:00
parent ebbd49cce5
commit c4c5c49239
Signed by: nyancat
GPG key ID: 91F97D9ED12639CF

View file

@ -1,13 +1,17 @@
package main package main
import ( import (
"context"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"io/fs"
"log" "log"
"net" "net"
"net/http" "net/http"
"os" "os"
"os/signal"
"strconv"
"strings" "strings"
"sync" "sync"
@ -27,10 +31,11 @@ const (
) )
type serverConfig struct { type serverConfig struct {
ASNDB string ASNDB string
CityDB string CityDB string
Listen string Listen string
Mode string Mode string
SocketPermission fs.FileMode
} }
type body struct { type body struct {
@ -47,11 +52,17 @@ type bodyError struct {
} }
func main() { func main() {
perm, err := strconv.ParseUint(getEnvOr("IP_CHECKER_SOCKET_PERMISSION", "0600"), 8, 32)
if err != nil {
log.Fatal(err)
}
config = serverConfig{ config = serverConfig{
ASNDB: os.Getenv("IP_CHECKER_ASN_DB"), ASNDB: os.Getenv("IP_CHECKER_ASN_DB"),
CityDB: os.Getenv("IP_CHECKER_CITY_DB"), CityDB: os.Getenv("IP_CHECKER_CITY_DB"),
Listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"), Listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"),
Mode: os.Getenv("IP_CHECKER_MODE"), SocketPermission: fs.FileMode(perm),
Mode: os.Getenv("IP_CHECKER_MODE"),
} }
if config.CityDB != "" { if config.CityDB != "" {
@ -76,8 +87,40 @@ func main() {
} }
http.HandleFunc("/api/v1/ip", handleRequest) http.HandleFunc("/api/v1/ip", handleRequest)
log.Printf("Starting server on %s", config.Listen) log.Printf("Starting server on %s", config.Listen)
log.Fatal(http.ListenAndServe(config.Listen, nil))
var srv http.Server
// graceful shutdown
go func() {
sigint := make(chan os.Signal, 1)
signal.Notify(sigint, os.Interrupt)
<-sigint
srv.Shutdown(context.Background())
}()
if strings.HasPrefix(config.Listen, "unix/") {
socketPath := strings.TrimPrefix(config.Listen, "unix/")
unixListener, err := net.Listen("unix", socketPath)
if err != nil {
log.Fatal(err)
}
if err := os.Chmod(socketPath, config.SocketPermission); err != nil {
log.Fatal(err)
}
if err := srv.Serve(unixListener); err != http.ErrServerClosed {
log.Fatal(err)
}
} else {
srv.Addr = config.Listen
if err := srv.ListenAndServe(); err != http.ErrServerClosed {
log.Fatal(err)
}
}
} }
func generateJSON(ip string) (body, error) { func generateJSON(ip string) (body, error) {