From 49ffcadb0a8a8378096af2f90ce40c233d877a04 Mon Sep 17 00:00:00 2001 From: dcarrillo Date: Sun, 19 Jul 2026 12:09:52 +0200 Subject: [PATCH] Fix data races, nil-record panics, scanner SSRF, log date layout, shutdown timeouts, and enable stricter linters (#54) High (behavioral bugs): - service.Geo: add RLock on lookups, handle Reload error, keep old DBs on failure - models.GeoDB.Reload: open new DBs before closing old ones (no reader gap) - router: guard nil geo records (404 instead of 500) in geo.go, generic.go, dns.go - router/port_scanner: reject undetermined client IPs (nil->localhost dial) - internal/httputils: fix access log date layout (02/Jan/2006, not literal Nov) - server/quic: log.Printf instead of log.Fatal in request handler Medium: - server/*: 10s shutdown timeout; SIGTERM drains servers before closing geo DBs - httputils: Clone() request headers instead of mutating the live map; strings.Builder - setting/app: handle all os.Stat errors; fix "truster" typos; use errors.New; remove dead Token field - cmd: errors.Is for sentinel checks; os.Stderr; log.Fatal instead of panic; pass http.Handler by value (not pointer-to-interface) - resolver: validate config at startup (pre-parse RRs, validate IPs); check WriteMsg errors; Setup returns error Low: - .golangci.yaml: enable errcheck, govet, errorlint; fix all resulting findings - server/tls: explicit TLSConfig{MinVersion: tls.VersionTLS12} - router/dns: label-exact vhost suffix check, strip port; template.Must - models: errors.Join instead of %s on []error - service/geo_test: race-regression test (concurrent lookups during Reload) - README: Go >= 1.25; fix latent test bug (store.Add->Set in dns_test.go) --- .golangci.yaml | 3 ++ README.md | 2 +- cmd/whatismyip.go | 17 +++--- integration-tests/integration_test.go | 12 ++++- internal/httputils/http.go | 21 +++----- internal/httputils/http_test.go | 4 +- internal/setting/app.go | 19 ++++--- models/geo.go | 44 +++++++-------- models/geo_test.go | 2 +- resolver/setup.go | 77 ++++++++++++++++----------- router/dns.go | 23 +++++--- router/dns_test.go | 4 +- router/generic.go | 33 ++++++------ router/geo.go | 39 ++++++++++---- router/port_scanner.go | 10 +++- router/setup.go | 3 +- server/dns.go | 4 +- server/prometheus.go | 4 +- server/quic.go | 8 +-- server/server.go | 10 ++-- server/tcp.go | 10 ++-- server/tls.go | 14 +++-- service/geo.go | 19 ++++++- service/geo_test.go | 34 ++++++++++++ 24 files changed, 273 insertions(+), 143 deletions(-) diff --git a/.golangci.yaml b/.golangci.yaml index d4d74e2..d09ba1e 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -2,6 +2,9 @@ version: "2" linters: default: none enable: + - errcheck + - errorlint + - govet - ineffassign - nakedret - revive diff --git a/README.md b/README.md index 61ce0f1..2456982 100644 --- a/README.md +++ b/README.md @@ -120,7 +120,7 @@ curl $(cat /proc/sys/kernel/random/uuid).dns.ifconfig.es ## Build -Golang >= 1.24 is required. +Golang >= 1.25 is required. `make build` diff --git a/cmd/whatismyip.go b/cmd/whatismyip.go index e053aef..74d0e49 100644 --- a/cmd/whatismyip.go +++ b/cmd/whatismyip.go @@ -2,8 +2,10 @@ package main import ( "context" + "errors" "flag" "fmt" + "log" "net/http" "os" "slices" @@ -25,11 +27,11 @@ import ( func main() { o, err := setting.Setup(os.Args[1:]) if err != nil { - if err == flag.ErrHelp || err == setting.ErrVersion { + if errors.Is(err, flag.ErrHelp) || errors.Is(err, setting.ErrVersion) { fmt.Print(o) os.Exit(0) } - fmt.Println(err) + fmt.Fprintln(os.Stderr, err) os.Exit(1) } @@ -38,7 +40,10 @@ func main() { if setting.App.Resolver.Domain != "" { store := cache.New(1*time.Minute, 10*time.Minute) - dnsEngine := resolver.Setup(store) + var dnsEngine *resolver.Resolver + if dnsEngine, err = resolver.Setup(store); err != nil { + log.Fatalf("Invalid resolver configuration: %s", err) + } nameServer := server.NewDNSServer(context.Background(), dnsEngine.Handler()) servers = append(servers, nameServer) engine.Use(router.GetDNSDiscoveryHandler(store, setting.App.Resolver.Domain, setting.App.Resolver.RedirectPort)) @@ -47,7 +52,7 @@ func main() { var geoSvc *service.Geo if setting.App.GeodbPath.City != "" || setting.App.GeodbPath.ASN != "" { if geoSvc, err = service.NewGeo(context.Background(), setting.App.GeodbPath.City, setting.App.GeodbPath.ASN); err != nil { - panic(err) + log.Fatalf("Failed to load geo databases: %s", err) } } @@ -92,12 +97,12 @@ func setupHTTPServers(ctx context.Context, handler http.Handler) []server.Server var servers []server.Server if setting.App.BindAddress != "" { - tcpServer := server.NewTCPServer(ctx, &handler) + tcpServer := server.NewTCPServer(ctx, handler) servers = append(servers, tcpServer) } if setting.App.TLSAddress != "" { - tlsServer := server.NewTLSServer(ctx, &handler) + tlsServer := server.NewTLSServer(ctx, handler) servers = append(servers, tlsServer) if setting.App.EnableHTTP3 { quicServer := server.NewQuicServer(ctx, tlsServer) diff --git a/integration-tests/integration_test.go b/integration-tests/integration_test.go index 9102939..4d0404f 100644 --- a/integration-tests/integration_test.go +++ b/integration-tests/integration_test.go @@ -142,7 +142,11 @@ func TestContainerIntegration(t *testing.T) { Started: true, }) require.NoError(t, err) - t.Cleanup(func() { c.Terminate(ctx) }) + t.Cleanup(func() { + if err := c.Terminate(ctx); err != nil { + t.Logf("Failed to terminate container: %s", err) + } + }) http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: true} tests := []struct { @@ -301,7 +305,11 @@ func TestContainerIntegrationDisableScan(t *testing.T) { Started: true, }) require.NoError(t, err) - t.Cleanup(func() { c.Terminate(ctx) }) + t.Cleanup(func() { + if err := c.Terminate(ctx); err != nil { + t.Logf("Failed to terminate container: %s", err) + } + }) t.Run("RequestScanEndpointWithDisabledScan", func(t *testing.T) { req, err := http.NewRequest("GET", "http://localhost:8000/scan/tcp/8000", nil) diff --git a/internal/httputils/http.go b/internal/httputils/http.go index 9b7be09..746bc58 100644 --- a/internal/httputils/http.go +++ b/internal/httputils/http.go @@ -11,32 +11,27 @@ import ( "github.com/gin-gonic/gin" ) -// HeadersToSortedString shorts and dumps http.Header to a string separated by \n +// HeadersToSortedString sorts and dumps http.Header to a string separated by \n func HeadersToSortedString(headers http.Header) string { - var output string - keys := make([]string, 0, len(headers)) for k := range headers { keys = append(keys, k) } sort.Strings(keys) + var output strings.Builder for _, k := range keys { - if len(headers[k]) > 1 { - for _, h := range headers[k] { - output += k + ": " + h + "\n" - } - } else { - output += k + ": " + headers[k][0] + "\n" + for _, h := range headers[k] { + output.WriteString(k + ": " + h + "\n") } } - return output + return output.String() } -// GetHeadersWithoutTrustedHeaders return a http.Heade object with the original headers except trusted headers +// GetHeadersWithoutTrustedHeaders returns a copy of the request headers with the trusted headers removed func GetHeadersWithoutTrustedHeaders(ctx *gin.Context) http.Header { - h := ctx.Request.Header + h := ctx.Request.Header.Clone() for _, k := range []string{setting.App.TrustedHeader, setting.App.TrustedPortHeader} { delete(h, textproto.CanonicalMIMEHeaderKey(k)) @@ -49,7 +44,7 @@ func GetHeadersWithoutTrustedHeaders(ctx *gin.Context) http.Header { func GetLogFormatter(param gin.LogFormatterParams) string { return fmt.Sprintf("%s - [%s] \"%s %s %s\" %d %d %d %s \"%s\" \"%s\" \"%s\"\n", param.ClientIP, - param.TimeStamp.Format("02/Nov/2006:15:04:05 -0700"), + param.TimeStamp.Format("02/Jan/2006:15:04:05 -0700"), param.Method, param.Path, param.Request.Proto, diff --git a/internal/httputils/http_test.go b/internal/httputils/http_test.go index aca1eb7..a4d7a37 100644 --- a/internal/httputils/http_test.go +++ b/internal/httputils/http_test.go @@ -25,7 +25,7 @@ Header3: Three } func TestGetLogFormatter(t *testing.T) { - expected := "127.0.0.1 - [01/Nov/0001:00:00:00 +0000] \"GET / HTTP/1.1\" 200 100 1000 local \"golang test 1.0\" \"1.1.1.1, 2.2.2.2\" \"-\"\n" + expected := "127.0.0.1 - [15/Jul/2022:10:30:00 +0000] \"GET / HTTP/1.1\" 200 100 1000 local \"golang test 1.0\" \"1.1.1.1, 2.2.2.2\" \"-\"\n" h := http.Header{} h.Set("User-Agent", "golang test 1.0") @@ -39,7 +39,7 @@ func TestGetLogFormatter(t *testing.T) { p := gin.LogFormatterParams{ ClientIP: "127.0.0.1", - TimeStamp: time.Time{}, + TimeStamp: time.Date(2022, time.July, 15, 10, 30, 0, 0, time.UTC), Method: "GET", Path: "/", StatusCode: 200, diff --git a/internal/setting/app.go b/internal/setting/app.go index 78b0373..9035c16 100644 --- a/internal/setting/app.go +++ b/internal/setting/app.go @@ -13,9 +13,8 @@ import ( ) type geodbConf struct { - City string - ASN string - Token *string + City string + ASN string } type serverSettings struct { ReadTimeout time.Duration @@ -136,25 +135,25 @@ func Setup(args []string) (output string, err error) { } if (App.GeodbPath.City != "" && App.GeodbPath.ASN == "") || (App.GeodbPath.City == "" && App.GeodbPath.ASN != "") { - return "", fmt.Errorf("both --geoip2-city and --geoip2-asn are mandatory to enable geo information") + return "", errors.New("both --geoip2-city and --geoip2-asn are mandatory to enable geo information") } if App.TrustedPortHeader != "" && App.TrustedHeader == "" { - return "", fmt.Errorf("truster-header is mandatory when truster-port-header is set") + return "", errors.New("trusted-header is mandatory when trusted-port-header is set") } if (App.TLSAddress != "") && (App.TLSCrtPath == "" || App.TLSKeyPath == "") { - return "", fmt.Errorf("in order to use TLS, the -tls-crt and -tls-key flags are mandatory") + return "", errors.New("in order to use TLS, the -tls-crt and -tls-key flags are mandatory") } if App.EnableHTTP3 && App.TLSAddress == "" { - return "", fmt.Errorf("in order to use HTTP3, the -tls-bind is mandatory") + return "", errors.New("in order to use HTTP3, the -tls-bind is mandatory") } if App.TemplatePath != "" { info, err := os.Stat(App.TemplatePath) - if os.IsNotExist(err) { - return "", fmt.Errorf("%s no such file or directory", App.TemplatePath) + if err != nil { + return "", fmt.Errorf("template path: %w", err) } if info.IsDir() { return "", fmt.Errorf("%s must be a file", App.TemplatePath) @@ -165,7 +164,7 @@ func Setup(args []string) (output string, err error) { var err error App.Resolver, err = readYAML(resolverConf) if err != nil { - return "", fmt.Errorf("error reading resolver configuration %w", err) + return "", fmt.Errorf("reading resolver configuration: %w", err) } } diff --git a/models/geo.go b/models/geo.go index 4afb7ab..882cea5 100644 --- a/models/geo.go +++ b/models/geo.go @@ -1,6 +1,7 @@ package models import ( + "errors" "fmt" "log" "net" @@ -53,39 +54,40 @@ func Setup(cityPath string, asnPath string) (*GeoDB, error) { } func (db *GeoDB) CloseDBs() error { - var errs []error + return closeReaders(db.City, db.ASN) +} - if db.City != nil { - if err := db.City.Close(); err != nil { - errs = append(errs, fmt.Errorf("closing city db: %w", err)) - } +func (db *GeoDB) Reload() error { + city, asn, err := openDatabases(db.cityPath, db.asnPath) + if err != nil { + return fmt.Errorf("opening new databases: %w", err) } - if db.ASN != nil { - if err := db.ASN.Close(); err != nil { - errs = append(errs, fmt.Errorf("closing ASN db: %w", err)) - } - } + oldCity, oldASN := db.City, db.ASN + db.City, db.ASN = city, asn - if len(errs) > 0 { - return fmt.Errorf("errors closing databases: %s", errs) + if err := closeReaders(oldCity, oldASN); err != nil { + return fmt.Errorf("closing previous databases: %w", err) } return nil } -func (db *GeoDB) Reload() error { - if err := db.CloseDBs(); err != nil { - return fmt.Errorf("closing existing connections: %w", err) +func closeReaders(city *maxminddb.Reader, asn *maxminddb.Reader) error { + var errs []error + + if city != nil { + if err := city.Close(); err != nil { + errs = append(errs, fmt.Errorf("closing city db: %w", err)) + } } - city, asn, err := openDatabases(db.cityPath, db.asnPath) - if err != nil { - return fmt.Errorf("opening new connections: %w", err) + if asn != nil { + if err := asn.Close(); err != nil { + errs = append(errs, fmt.Errorf("closing ASN db: %w", err)) + } } - db.City = city - db.ASN = asn - return nil + return errors.Join(errs...) } func (db *GeoDB) LookupCity(ip net.IP) (*GeoRecord, error) { diff --git a/models/geo_test.go b/models/geo_test.go index 7506f7e..166d5d0 100644 --- a/models/geo_test.go +++ b/models/geo_test.go @@ -63,7 +63,7 @@ func TestModels(t *testing.T) { db, err := Setup("../test/GeoIP2-City-Test.mmdb", "../test/GeoLite2-ASN-Test.mmdb") require.NoError(t, err, fmt.Sprintf("Error setting up db: %s", err)) - defer db.CloseDBs() + t.Cleanup(func() { assert.NoError(t, db.CloseDBs()) }) assert.NotNil(t, db.ASN) assert.NotNil(t, db.City) diff --git a/resolver/setup.go b/resolver/setup.go index 00beabf..a73bb96 100644 --- a/resolver/setup.go +++ b/resolver/setup.go @@ -1,6 +1,7 @@ package resolver import ( + "fmt" "log" "net" "strings" @@ -16,7 +17,7 @@ type Resolver struct { handler *dns.ServeMux store *cache.Cache domain string - rr []string + rr []dns.RR ipv4 []net.IP ipv6 []net.IP } @@ -28,27 +29,51 @@ func ensureDotSuffix(s string) string { return s } -func Setup(store *cache.Cache) *Resolver { - var ipv4, ipv6 []net.IP - for _, ip := range setting.App.Resolver.Ipv4 { - ipv4 = append(ipv4, net.ParseIP(ip)) +func Setup(store *cache.Cache) (*Resolver, error) { + domain := ensureDotSuffix(setting.App.Resolver.Domain) + + rr := make([]dns.RR, 0, len(setting.App.Resolver.ResourceRecords)) + for _, res := range setting.App.Resolver.ResourceRecords { + record, err := dns.NewRR(domain + " " + res) + if err != nil { + return nil, fmt.Errorf("parsing resource record %q: %w", res, err) + } + rr = append(rr, record) } - for _, ip := range setting.App.Resolver.Ipv6 { - ipv6 = append(ipv6, net.ParseIP(ip)) + + ipv4, err := parseIPs(setting.App.Resolver.Ipv4) + if err != nil { + return nil, fmt.Errorf("parsing ipv4 addresses: %w", err) + } + ipv6, err := parseIPs(setting.App.Resolver.Ipv6) + if err != nil { + return nil, fmt.Errorf("parsing ipv6 addresses: %w", err) } resolver := &Resolver{ handler: dns.NewServeMux(), store: store, - domain: ensureDotSuffix(setting.App.Resolver.Domain), - rr: setting.App.Resolver.ResourceRecords, + domain: domain, + rr: rr, ipv4: ipv4, ipv6: ipv6, } resolver.handler.HandleFunc(resolver.domain, resolver.resolve) resolver.handler.HandleFunc(".", resolver.blackHole) - return resolver + return resolver, nil +} + +func parseIPs(addresses []string) ([]net.IP, error) { + ips := make([]net.IP, 0, len(addresses)) + for _, address := range addresses { + ip := net.ParseIP(address) + if ip == nil { + return nil, fmt.Errorf("invalid IP address %q", address) + } + ips = append(ips, ip) + } + return ips, nil } func (rsv *Resolver) Handler() *dns.ServeMux { @@ -58,7 +83,7 @@ func (rsv *Resolver) Handler() *dns.ServeMux { func (rsv *Resolver) blackHole(w dns.ResponseWriter, r *dns.Msg) { msg := startReply(r) msg.SetRcode(r, dns.RcodeRefused) - w.WriteMsg(msg) + writeMsg(w, msg) logger(w, r.Question[0], msg.Rcode) metrics.RecordDNSQuery(dns.TypeToString[r.Question[0].Qtype], dns.RcodeToString[msg.Rcode]) } @@ -69,17 +94,10 @@ func (rsv *Resolver) resolve(w dns.ResponseWriter, r *dns.Msg) { ip, _, _ := net.SplitHostPort(w.RemoteAddr().String()) for _, res := range rsv.rr { - t := strings.Split(res, " ")[2] - if q.Qtype == dns.StringToType[t] { - brr, err := buildRR(rsv.domain + " " + res) - if err != nil { - msg.SetRcode(r, dns.RcodeServerFailure) - logger(w, q, msg.Rcode, err.Error()) - } else { - msg.Answer = append(msg.Answer, brr) - logger(w, q, msg.Rcode) - } - w.WriteMsg(msg) + if q.Qtype == res.Header().Rrtype { + msg.Answer = append(msg.Answer, dns.Copy(res)) + writeMsg(w, msg) + logger(w, q, msg.Rcode) metrics.RecordDNSQuery(dns.TypeToString[q.Qtype], dns.RcodeToString[msg.Rcode]) return } @@ -90,14 +108,16 @@ func (rsv *Resolver) resolve(w dns.ResponseWriter, r *dns.Msg) { switch { case uuid.IsValid(subDomain): msg.SetRcode(r, rsv.getIP(q, msg)) - rsv.store.Add(subDomain, ip, cache.DefaultExpiration) + // Add fails when the uuid is already registered; keep the first seen + // resolver IP for the discovery window + _ = rsv.store.Add(subDomain, ip, cache.DefaultExpiration) case lowerName == rsv.domain: msg.SetRcode(r, rsv.getIP(q, msg)) default: msg.SetRcode(r, dns.RcodeRefused) } - w.WriteMsg(msg) + writeMsg(w, msg) logger(w, q, msg.Rcode) metrics.RecordDNSQuery(dns.TypeToString[q.Qtype], dns.RcodeToString[msg.Rcode]) } @@ -126,13 +146,10 @@ func (rsv *Resolver) getIP(question dns.Question, msg *dns.Msg) int { return dns.RcodeRefused } -func buildRR(rrs string) (dns.RR, error) { - rr, err := dns.NewRR(rrs) - if err != nil { - return nil, err +func writeMsg(w dns.ResponseWriter, msg *dns.Msg) { + if err := w.WriteMsg(msg); err != nil { + log.Printf("Failed to write DNS response: %s", err) } - - return rr, nil } func setHdr(q dns.Question) dns.RR_Header { diff --git a/router/dns.go b/router/dns.go index a3e5fdd..cedad18 100644 --- a/router/dns.go +++ b/router/dns.go @@ -29,12 +29,13 @@ type dnsData struct { // Implement a proper vhost manager instead of using a middleware func GetDNSDiscoveryHandler(store *cache.Cache, domain string, redirectPort string) gin.HandlerFunc { return func(ctx *gin.Context) { - if !strings.HasSuffix(ctx.Request.Host, domain) { + host := hostWithoutPort(ctx.Request.Host) + if host != domain && !strings.HasSuffix(host, "."+domain) { ctx.Next() return } - if ctx.Request.Host == domain && ctx.Request.URL.Path == "/" { + if host == domain && ctx.Request.URL.Path == "/" { ctx.Redirect(http.StatusFound, fmt.Sprintf("http://%s.%s%s", uuid.New().String(), domain, redirectPort)) ctx.Abort() return @@ -45,6 +46,13 @@ func GetDNSDiscoveryHandler(store *cache.Cache, domain string, redirectPort stri } } +func hostWithoutPort(host string) string { + if h, _, err := net.SplitHostPort(host); err == nil { + return h + } + return host +} + func handleDNS(ctx *gin.Context, store *cache.Cache) { d := strings.Split(ctx.Request.Host, ".")[0] if !validator.IsValid(d) { @@ -72,12 +80,11 @@ func handleDNS(ctx *gin.Context, store *cache.Cache) { geoResp := dnsGeoData{} if geoSvc != nil { - cityRecord := geoSvc.LookUpCity(ip) - asnRecord := geoSvc.LookUpASN(ip) - - geoResp = dnsGeoData{ - Country: cityRecord.Country.Names["en"], - AsnOrganization: asnRecord.AutonomousSystemOrganization, + if cityRecord := geoSvc.LookUpCity(ip); cityRecord != nil { + geoResp.Country = cityRecord.Country.Names["en"] + } + if asnRecord := geoSvc.LookUpASN(ip); asnRecord != nil { + geoResp.AsnOrganization = asnRecord.AutonomousSystemOrganization } } diff --git a/router/dns_test.go b/router/dns_test.go index a2e45ba..5688516 100644 --- a/router/dns_test.go +++ b/router/dns_test.go @@ -100,7 +100,7 @@ func TestHandleDNS(t *testing.T) { req.Host = tt.subDomain + "." + domain if tt.stored != "" { - store.Add(tt.subDomain, tt.stored, cache.DefaultExpiration) + store.Set(tt.subDomain, tt.stored, cache.DefaultExpiration) } w := httptest.NewRecorder() @@ -143,7 +143,7 @@ func TestAcceptDNSRequest(t *testing.T) { c, _ := gin.CreateTestContext(w) c.Request = req - store.Add(u, testIP.ipv4, cache.DefaultExpiration) + store.Set(u, testIP.ipv4, cache.DefaultExpiration) handleDNS(c, store) assert.Equal(t, http.StatusOK, w.Code) diff --git a/router/generic.go b/router/generic.go index e374a1c..7ce4459 100644 --- a/router/generic.go +++ b/router/generic.go @@ -75,8 +75,12 @@ func getAllAsString(ctx *gin.Context) { output += "Client Port: " + getClientPort(ctx) + "\n" if geoSvc != nil { - output += geoCityRecordToString(geoSvc.LookUpCity(ip)) + "\n" - output += geoASNRecordToString(geoSvc.LookUpASN(ip)) + "\n" + if cityRecord := geoSvc.LookUpCity(ip); cityRecord != nil { + output += geoCityRecordToString(cityRecord) + "\n" + } + if asnRecord := geoSvc.LookUpASN(ip); asnRecord != nil { + output += geoASNRecordToString(asnRecord) + "\n" + } } h := httputils.GetHeadersWithoutTrustedHeaders(ctx) @@ -100,19 +104,18 @@ func jsonOutput(ctx *gin.Context) JSONResponse { geoResp := GeoResponse{} if geoSvc != nil { - cityRecord := geoSvc.LookUpCity(ip) - asnRecord := geoSvc.LookUpASN(ip) - - geoResp = GeoResponse{ - Country: cityRecord.Country.Names["en"], - CountryCode: cityRecord.Country.ISOCode, - City: cityRecord.City.Names["en"], - Latitude: cityRecord.Location.Latitude, - Longitude: cityRecord.Location.Longitude, - PostalCode: cityRecord.Postal.Code, - TimeZone: cityRecord.Location.TimeZone, - ASN: asnRecord.AutonomousSystemNumber, - ASNOrganization: asnRecord.AutonomousSystemOrganization, + if cityRecord := geoSvc.LookUpCity(ip); cityRecord != nil { + geoResp.Country = cityRecord.Country.Names["en"] + geoResp.CountryCode = cityRecord.Country.ISOCode + geoResp.City = cityRecord.City.Names["en"] + geoResp.Latitude = cityRecord.Location.Latitude + geoResp.Longitude = cityRecord.Location.Longitude + geoResp.PostalCode = cityRecord.Postal.Code + geoResp.TimeZone = cityRecord.Location.TimeZone + } + if asnRecord := geoSvc.LookUpASN(ip); asnRecord != nil { + geoResp.ASN = asnRecord.AutonomousSystemNumber + geoResp.ASNOrganization = asnRecord.AutonomousSystemOrganization } } diff --git a/router/geo.go b/router/geo.go index 06ec020..5a1f97c 100644 --- a/router/geo.go +++ b/router/geo.go @@ -87,15 +87,24 @@ func getGeoAsString(ctx *gin.Context) { return } - field := strings.ToLower(ctx.Params.ByName("field")) record := geoSvc.LookUpCity(net.ParseIP(ctx.ClientIP())) + if record == nil { + ctx.String(http.StatusNotFound, http.StatusText(http.StatusNotFound)) + return + } + + field := strings.ToLower(ctx.Params.ByName("field")) if field == "" { ctx.String(http.StatusOK, geoCityRecordToString(record)) - } else if g, ok := geoOutput[field]; ok { - ctx.String(http.StatusOK, g.format(record)) - } else { - ctx.String(http.StatusNotFound, http.StatusText(http.StatusNotFound)) + return } + + g, ok := geoOutput[field] + if !ok { + ctx.String(http.StatusNotFound, http.StatusText(http.StatusNotFound)) + return + } + ctx.String(http.StatusOK, g.format(record)) } func getASNAsString(ctx *gin.Context) { @@ -103,15 +112,25 @@ func getASNAsString(ctx *gin.Context) { ctx.String(http.StatusNotFound, http.StatusText(http.StatusNotFound)) return } - field := strings.ToLower(ctx.Params.ByName("field")) + record := geoSvc.LookUpASN(net.ParseIP(ctx.ClientIP())) + if record == nil { + ctx.String(http.StatusNotFound, http.StatusText(http.StatusNotFound)) + return + } + + field := strings.ToLower(ctx.Params.ByName("field")) if field == "" { ctx.String(http.StatusOK, geoASNRecordToString(record)) - } else if g, ok := asnOutput[field]; ok { - ctx.String(http.StatusOK, g.format(record)) - } else { - ctx.String(http.StatusNotFound, http.StatusText(http.StatusNotFound)) + return } + + g, ok := asnOutput[field] + if !ok { + ctx.String(http.StatusNotFound, http.StatusText(http.StatusNotFound)) + return + } + ctx.String(http.StatusOK, g.format(record)) } func geoCityRecordToString(record *models.GeoRecord) string { diff --git a/router/port_scanner.go b/router/port_scanner.go index f9d7562..c6f604b 100644 --- a/router/port_scanner.go +++ b/router/port_scanner.go @@ -29,8 +29,16 @@ func scanTCPPort(ctx *gin.Context) { return } + ip := net.ParseIP(ctx.ClientIP()) + if ip == nil { + ctx.JSON(http.StatusBadRequest, JSONScanResponse{ + Reason: "client ip could not be determined", + }) + return + } + add := net.TCPAddr{ - IP: net.ParseIP(ctx.ClientIP()), + IP: ip, Port: port, } diff --git a/router/setup.go b/router/setup.go index ded52c6..cd65694 100644 --- a/router/setup.go +++ b/router/setup.go @@ -13,8 +13,7 @@ var geoSvc *service.Geo func SetupTemplate(r *gin.Engine) { if setting.App.TemplatePath == "" { - t, _ := template.New("home").Parse(home) - r.SetHTMLTemplate(t) + r.SetHTMLTemplate(template.Must(template.New("home").Parse(home))) } else { log.Printf("Template %s has been loaded", setting.App.TemplatePath) r.LoadHTMLFiles(setting.App.TemplatePath) diff --git a/server/dns.go b/server/dns.go index a79b102..e30ea42 100644 --- a/server/dns.go +++ b/server/dns.go @@ -42,7 +42,9 @@ func (d *DNS) Start() { func (d *DNS) Stop() { log.Print("Stopping DNS server...") - if err := d.server.Shutdown(); err != nil { + ctx, cancel := context.WithTimeout(d.ctx, shutdownTimeout) + defer cancel() + if err := d.server.ShutdownContext(ctx); err != nil { log.Printf("DNS server forced to shutdown: %s", err) } } diff --git a/server/prometheus.go b/server/prometheus.go index 50848c1..88990c5 100644 --- a/server/prometheus.go +++ b/server/prometheus.go @@ -42,7 +42,9 @@ func (p *Prometheus) Start() { func (p *Prometheus) Stop() { log.Print("Stopping Prometheus server...") - if err := p.server.Shutdown(p.ctx); err != nil { + ctx, cancel := context.WithTimeout(p.ctx, shutdownTimeout) + defer cancel() + if err := p.server.Shutdown(ctx); err != nil { log.Printf("Prometheus server forced to shutdown: %s", err) } } diff --git a/server/quic.go b/server/quic.go index a057264..877204f 100644 --- a/server/quic.go +++ b/server/quic.go @@ -32,7 +32,7 @@ func (q *Quic) Start() { parentHandler := q.tlsServer.server.Handler q.tlsServer.server.Handler = http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { if err := q.server.SetQUICHeaders(rw.Header()); err != nil { - log.Fatal(err) + log.Printf("Failed to set QUIC headers: %s", err) } parentHandler.ServeHTTP(rw, req) @@ -49,7 +49,9 @@ func (q *Quic) Start() { func (q *Quic) Stop() { log.Print("Stopping QUIC server...") - if err := q.server.Close(); err != nil { - log.Print("QUIC server forced to shutdown") + ctx, cancel := context.WithTimeout(q.ctx, shutdownTimeout) + defer cancel() + if err := q.server.Shutdown(ctx); err != nil { + log.Printf("QUIC server forced to shutdown: %s", err) } } diff --git a/server/server.go b/server/server.go index d2a7b09..eb828c4 100644 --- a/server/server.go +++ b/server/server.go @@ -5,10 +5,13 @@ import ( "os" "os/signal" "syscall" + "time" "github.com/dcarrillo/whatismyip/service" ) +const shutdownTimeout = 10 * time.Second + type Server interface { Start() Stop() @@ -29,11 +32,10 @@ func Setup(servers []Server, geoSvc *service.Geo) *Manager { func (m *Manager) Run() { m.start() - signalChan := make(chan os.Signal, len(m.servers)) + signalChan := make(chan os.Signal, 1) signal.Notify(signalChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM) - var s os.Signal for { - s = <-signalChan + s := <-signalChan if s == syscall.SIGHUP { m.stop() @@ -43,10 +45,10 @@ func (m *Manager) Run() { m.start() } else { log.Print("Shutting down...") + m.stop() if m.geoSvc != nil { m.geoSvc.Shutdown() } - m.stop() break } } diff --git a/server/tcp.go b/server/tcp.go index 6df3125..35075b6 100644 --- a/server/tcp.go +++ b/server/tcp.go @@ -11,11 +11,11 @@ import ( type TCP struct { server *http.Server - handler *http.Handler + handler http.Handler ctx context.Context } -func NewTCPServer(ctx context.Context, handler *http.Handler) *TCP { +func NewTCPServer(ctx context.Context, handler http.Handler) *TCP { return &TCP{ handler: handler, ctx: ctx, @@ -25,7 +25,7 @@ func NewTCPServer(ctx context.Context, handler *http.Handler) *TCP { func (t *TCP) Start() { t.server = &http.Server{ Addr: setting.App.BindAddress, - Handler: *t.handler, + Handler: t.handler, ReadTimeout: setting.App.Server.ReadTimeout, WriteTimeout: setting.App.Server.WriteTimeout, } @@ -40,7 +40,9 @@ func (t *TCP) Start() { func (t *TCP) Stop() { log.Print("Stopping TCP server...") - if err := t.server.Shutdown(t.ctx); err != nil { + ctx, cancel := context.WithTimeout(t.ctx, shutdownTimeout) + defer cancel() + if err := t.server.Shutdown(ctx); err != nil { log.Printf("TCP server forced to shutdown: %s", err) } } diff --git a/server/tls.go b/server/tls.go index 15ef776..795116a 100644 --- a/server/tls.go +++ b/server/tls.go @@ -2,6 +2,7 @@ package server import ( "context" + "crypto/tls" "errors" "log" "net/http" @@ -11,11 +12,11 @@ import ( type TLS struct { server *http.Server - handler *http.Handler + handler http.Handler ctx context.Context } -func NewTLSServer(ctx context.Context, handler *http.Handler) *TLS { +func NewTLSServer(ctx context.Context, handler http.Handler) *TLS { return &TLS{ handler: handler, ctx: ctx, @@ -25,9 +26,12 @@ func NewTLSServer(ctx context.Context, handler *http.Handler) *TLS { func (t *TLS) Start() { t.server = &http.Server{ Addr: setting.App.TLSAddress, - Handler: *t.handler, + Handler: t.handler, ReadTimeout: setting.App.Server.ReadTimeout, WriteTimeout: setting.App.Server.WriteTimeout, + TLSConfig: &tls.Config{ + MinVersion: tls.VersionTLS12, + }, } log.Printf("Starting TLS server listening on %s", setting.App.TLSAddress) @@ -41,7 +45,9 @@ func (t *TLS) Start() { func (t *TLS) Stop() { log.Print("Stopping TLS server...") - if err := t.server.Shutdown(t.ctx); err != nil { + ctx, cancel := context.WithTimeout(t.ctx, shutdownTimeout) + defer cancel() + if err := t.server.Shutdown(ctx); err != nil { log.Printf("TLS server forced to shutdown: %s", err) } } diff --git a/service/geo.go b/service/geo.go index a734b5d..3282a9d 100644 --- a/service/geo.go +++ b/service/geo.go @@ -36,6 +36,9 @@ func NewGeo(ctx context.Context, cityPath string, asnPath string) (*Geo, error) } func (g *Geo) LookUpCity(ip net.IP) *models.GeoRecord { + g.mu.RLock() + defer g.mu.RUnlock() + record, err := g.db.LookupCity(ip) if err != nil { log.Print(err) @@ -47,6 +50,9 @@ func (g *Geo) LookUpCity(ip net.IP) *models.GeoRecord { } func (g *Geo) LookUpASN(ip net.IP) *models.ASNRecord { + g.mu.RLock() + defer g.mu.RUnlock() + record, err := g.db.LookupASN(ip) if err != nil { log.Print(err) @@ -59,7 +65,13 @@ func (g *Geo) LookUpASN(ip net.IP) *models.ASNRecord { func (g *Geo) Shutdown() { g.cancel() - g.db.CloseDBs() + + g.mu.Lock() + defer g.mu.Unlock() + + if err := g.db.CloseDBs(); err != nil { + log.Printf("Error closing geo databases: %s", err) + } } func (g *Geo) Reload() { @@ -71,6 +83,9 @@ func (g *Geo) Reload() { g.mu.Lock() defer g.mu.Unlock() - g.db.Reload() + if err := g.db.Reload(); err != nil { + log.Printf("Geo database reload failed: %s", err) + return + } log.Print("Geo database reloaded") } diff --git a/service/geo_test.go b/service/geo_test.go index 25508cb..98909fc 100644 --- a/service/geo_test.go +++ b/service/geo_test.go @@ -4,9 +4,11 @@ import ( "context" "net" "os" + "sync" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) var geoSvc *Geo @@ -31,3 +33,35 @@ func TestASNLookup(t *testing.T) { a = geoSvc.LookUpASN(net.ParseIP("1.1.1.1")) assert.NotNil(t, a) } + +func TestReloadIsSafeForConcurrentLookups(t *testing.T) { + svc, err := NewGeo(context.Background(), "../test/GeoIP2-City-Test.mmdb", "../test/GeoLite2-ASN-Test.mmdb") + require.NoError(t, err) + defer svc.Shutdown() + + ip := net.ParseIP("81.2.69.192") + + var wg sync.WaitGroup + stop := make(chan struct{}) + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + svc.LookUpCity(ip) + svc.LookUpASN(ip) + } + } + }() + + for range 100 { + svc.Reload() + } + close(stop) + wg.Wait() + + assert.NotNil(t, svc.LookUpCity(ip)) +}