diff --git a/api.go b/api.go index 960d99f2..5e057ded 100644 --- a/api.go +++ b/api.go @@ -1,6 +1,7 @@ package main import ( + "context" "encoding/json" "fmt" "html/template" @@ -29,7 +30,7 @@ const cacheScanTime = time.Minute // Type for performing checks against an input domain. Returns // a DomainResult object from the checker. -type checkPerformer func(API, string) (checker.DomainResult, error) +type checkPerformer func(context.Context, API, string) (checker.DomainResult, error) // API is the HTTP API that this service provides. // All requests respond with an APIResponse JSON, with fields: @@ -89,7 +90,7 @@ func (api *API) wrapper(handler apiHandler) func(w http.ResponseWriter, r *http. } } -func defaultCheck(api API, domain string) (checker.DomainResult, error) { +func defaultCheck(ctx context.Context, api API, domain string) (checker.DomainResult, error) { policyChan := models.Domain{Name: domain}.AsyncPolicyListCheck(api.Database, api.List) c := checker.Checker{ Cache: &checker.ScanCache{ @@ -98,7 +99,7 @@ func defaultCheck(api API, domain string) (checker.DomainResult, error) { }, Timeout: 3 * time.Second, } - result := c.CheckDomain(domain, nil) + result := c.CheckDomain(ctx, domain, nil) policyResult := <-policyChan result.ExtraResults["policylist"] = &policyResult return result, nil @@ -135,7 +136,7 @@ func (api API) Scan(r *http.Request) APIResponse { } } // 1. Conduct scan via starttls-checker - scanData, err := api.CheckDomain(api, domain) + scanData, err := api.CheckDomain(r.Context(), api, domain) if err != nil { return APIResponse{StatusCode: http.StatusInternalServerError, Message: err.Error()} } diff --git a/checker/checker.go b/checker/checker.go index 4d7d619e..3a4379db 100644 --- a/checker/checker.go +++ b/checker/checker.go @@ -1,6 +1,7 @@ package checker import ( + "context" "net" "time" ) @@ -22,7 +23,7 @@ type Checker struct { // CheckHostname defines the function that should be used to check each hostname. // If nil, FullCheckHostname (all hostname checks) will be used. - CheckHostname func(string, string, time.Duration) HostnameResult + CheckHostname func(context.Context, string, string, time.Duration) HostnameResult // checkMTASTSOverride is used to mock MTA-STS checks. checkMTASTSOverride func(string, map[string]HostnameResult) *MTASTSResult diff --git a/checker/cmd/starttls-check/cmd.go b/checker/cmd/starttls-check/cmd.go index 10d83f7e..495b7b29 100644 --- a/checker/cmd/starttls-check/cmd.go +++ b/checker/cmd/starttls-check/cmd.go @@ -2,6 +2,7 @@ package main import ( "bufio" + "context" "encoding/csv" "encoding/json" "flag" @@ -55,7 +56,7 @@ func main() { if *domain != "" { // Handle single domain and return - result := c.CheckDomain(*domain, nil) + result := c.CheckDomain(context.Background(), *domain, nil) resultHandler.HandleDomain(result) os.Exit(0) } diff --git a/checker/domain.go b/checker/domain.go index 8203276d..a9bf9a94 100644 --- a/checker/domain.go +++ b/checker/domain.go @@ -63,15 +63,15 @@ func (d DomainResult) setStatus(status DomainStatus) DomainResult { return d } -func lookupMXWithTimeout(domain string, timeout time.Duration) ([]*net.MX, error) { - ctx, cancel := context.WithTimeout(context.TODO(), timeout) +func lookupMXWithTimeout(ctx context.Context, domain string, timeout time.Duration) ([]*net.MX, error) { + ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() var r net.Resolver return r.LookupMX(ctx, domain) } // lookupHostnames retrieves the MX hostnames associated with a domain. -func (c *Checker) lookupHostnames(domain string) ([]string, error) { +func (c *Checker) lookupHostnames(ctx context.Context, domain string) ([]string, error) { domainASCII, err := idna.ToASCII(domain) if err != nil { return nil, fmt.Errorf("domain name %s couldn't be converted to ASCII", domain) @@ -81,7 +81,7 @@ func (c *Checker) lookupHostnames(domain string) ([]string, error) { if c.lookupMXOverride != nil { mxs, err = c.lookupMXOverride(domain) } else { - mxs, err = lookupMXWithTimeout(domainASCII, c.timeout()) + mxs, err = lookupMXWithTimeout(ctx, domainASCII, c.timeout()) } if err != nil || len(mxs) == 0 { return nil, fmt.Errorf("No MX records found") @@ -104,7 +104,7 @@ func (c *Checker) lookupHostnames(domain string) ([]string, error) { // `domain` is the mail domain to perform the lookup on. // `expectedHostnames` is the list of expected hostnames. // If `expectedHostnames` is nil, we don't validate the DNS lookup. -func (c *Checker) CheckDomain(domain string, expectedHostnames []string) DomainResult { +func (c *Checker) CheckDomain(ctx context.Context, domain string, expectedHostnames []string) DomainResult { result := DomainResult{ Domain: domain, MxHostnames: expectedHostnames, @@ -114,20 +114,20 @@ func (c *Checker) CheckDomain(domain string, expectedHostnames []string) DomainR // 1. Look up hostnames // 2. Perform and aggregate checks from those hostnames. // 3. Set a summary message. - hostnames, err := c.lookupHostnames(domain) + hostnames, err := c.lookupHostnames(ctx, domain) if err != nil { return result.setStatus(DomainCouldNotConnect) } checkedHostnames := make([]string, 0) for _, hostname := range hostnames { - hostnameResult := c.checkHostname(domain, hostname) + hostnameResult := c.checkHostname(ctx, domain, hostname) result.HostnameResults[hostname] = hostnameResult if hostnameResult.couldConnect() { checkedHostnames = append(checkedHostnames, hostname) } } result.PreferredHostnames = checkedHostnames - result.MTASTSResult = c.checkMTASTS(domain, result.HostnameResults) + result.MTASTSResult = c.checkMTASTS(ctx, domain, result.HostnameResults) // Derive Domain code from Hostname results. if len(checkedHostnames) == 0 { diff --git a/checker/domain_test.go b/checker/domain_test.go index c29d5a73..8d334c83 100644 --- a/checker/domain_test.go +++ b/checker/domain_test.go @@ -1,6 +1,7 @@ package checker import ( + "context" "fmt" "net" "testing" @@ -59,7 +60,7 @@ func mockLookupMX(domain string) ([]*net.MX, error) { return result, nil } -func mockCheckHostname(domain string, hostname string, _ time.Duration) HostnameResult { +func mockCheckHostname(_ context.Context, domain string, hostname string, _ time.Duration) HostnameResult { if result, ok := hostnameResults[hostname]; ok { return HostnameResult{ Result: &result, @@ -120,7 +121,7 @@ func performTestsWithCacheTimeout(t *testing.T, tests []domainTestCase, cacheExp if test.expectedHostnames == nil { test.expectedHostnames = mxLookup[test.domain] } - got := c.CheckDomain(test.domain, test.expectedHostnames).Status + got := c.CheckDomain(context.Background(), test.domain, test.expectedHostnames).Status test.check(t, got) } } diff --git a/checker/hostname.go b/checker/hostname.go index 96fa5a03..25b8e4a4 100644 --- a/checker/hostname.go +++ b/checker/hostname.go @@ -1,6 +1,7 @@ package checker import ( + "context" "crypto/tls" "crypto/x509" "net" @@ -70,11 +71,13 @@ func getThisHostname() string { // Performs an SMTP dial with a short timeout. // https://github.com/golang/go/issues/16436 -func smtpDialWithTimeout(hostname string, timeout time.Duration) (*smtp.Client, error) { +func smtpDialWithTimeout(ctx context.Context, hostname string, timeout time.Duration) (*smtp.Client, error) { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() if _, _, err := net.SplitHostPort(hostname); err != nil { hostname += ":25" } - conn, err := net.DialTimeout("tcp", hostname, timeout) + conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", hostname) if err != nil { return nil, err } @@ -160,13 +163,13 @@ func tlsConfigForCipher(ciphers []uint16) tls.Config { } // Checks to see that insecure ciphers are disabled. -func checkTLSCipher(hostname string, timeout time.Duration) *Result { +func checkTLSCipher(ctx context.Context, hostname string, timeout time.Duration) *Result { result := MakeResult("cipher") badCiphers := []uint16{ tls.TLS_RSA_WITH_RC4_128_SHA, tls.TLS_ECDHE_ECDSA_WITH_RC4_128_SHA, tls.TLS_ECDHE_RSA_WITH_RC4_128_SHA} - client, err := smtpDialWithTimeout(hostname, timeout) + client, err := smtpDialWithTimeout(ctx, hostname, timeout) if err != nil { return result.Error("Could not establish connection with hostname %s", hostname) } @@ -179,7 +182,7 @@ func checkTLSCipher(hostname string, timeout time.Duration) *Result { return result.Success() } -func checkTLSVersion(client *smtp.Client, hostname string, timeout time.Duration) *Result { +func checkTLSVersion(ctx context.Context, client *smtp.Client, hostname string, timeout time.Duration) *Result { result := MakeResult(Version) // Check the TLS version of the existing connection. @@ -193,7 +196,7 @@ func checkTLSVersion(client *smtp.Client, hostname string, timeout time.Duration } // Attempt to connect with an old SSL version. - client, err := smtpDialWithTimeout(hostname, timeout) + client, err := smtpDialWithTimeout(ctx, hostname, timeout) if err != nil { return result.Error("Could not establish connection: %v", err) } @@ -212,7 +215,7 @@ func checkTLSVersion(client *smtp.Client, hostname string, timeout time.Duration // checkHostname returns the result of c.CheckHostname or FullCheckHostname, // using or updating the Checker's cache. -func (c *Checker) checkHostname(domain string, hostname string) HostnameResult { +func (c *Checker) checkHostname(ctx context.Context, domain string, hostname string) HostnameResult { check := c.CheckHostname if check == nil { // If CheckHostname hasn't been set, default to the full set of checks. @@ -220,18 +223,18 @@ func (c *Checker) checkHostname(domain string, hostname string) HostnameResult { } if c.Cache == nil { - return check(domain, hostname, c.timeout()) + return check(ctx, domain, hostname, c.timeout()) } hostnameResult, err := c.Cache.GetHostnameScan(hostname) if err != nil { - hostnameResult = check(domain, hostname, c.timeout()) + hostnameResult = check(ctx, domain, hostname, c.timeout()) c.Cache.PutHostnameScan(hostname, hostnameResult) } return hostnameResult } // NoopCheckHostname returns a fake error result containing `domain` and `hostname`. -func NoopCheckHostname(domain string, hostname string, _ time.Duration) HostnameResult { +func NoopCheckHostname(ctx context.Context, domain string, hostname string, _ time.Duration) HostnameResult { r := HostnameResult{ Domain: domain, Hostname: hostname, @@ -244,8 +247,8 @@ func NoopCheckHostname(domain string, hostname string, _ time.Duration) Hostname // FullCheckHostname performs a series of checks against a hostname for an email domain. // `domain` is the mail domain that this server serves email for. // `hostname` is the hostname for this server. -func FullCheckHostname(domain string, hostname string, timeout time.Duration) HostnameResult { - result := HostnameResult{ +func FullCheckHostname(ctx context.Context, domain string, hostname string, timeout time.Duration) HostnameResult { + result := &HostnameResult{ Domain: domain, Hostname: hostname, Result: MakeResult("hostnames"), @@ -254,23 +257,23 @@ func FullCheckHostname(domain string, hostname string, timeout time.Duration) Ho // Connect to the SMTP server and use that connection to perform as many checks as possible. connectivityResult := MakeResult(Connectivity) - client, err := smtpDialWithTimeout(hostname, timeout) + client, err := smtpDialWithTimeout(ctx, hostname, timeout) if err != nil { result.addCheck(connectivityResult.Error("Could not establish connection: %v", err)) - return result + return *result } - defer client.Close() + defer client.Quit() result.addCheck(connectivityResult.Success()) result.addCheck(checkStartTLS(client)) if result.Status != Success { - return result + return *result } result.addCheck(checkCert(client, domain, hostname)) // result.addCheck(checkTLSCipher(hostname)) // Creates a new connection to check for SSLv2/3 support because we can't call starttls twice. - result.addCheck(checkTLSVersion(client, hostname, timeout)) + result.addCheck(checkTLSVersion(ctx, client, hostname, timeout)) - return result + return *result } diff --git a/checker/hostname_test.go b/checker/hostname_test.go index 29a38fd9..2a9f0992 100644 --- a/checker/hostname_test.go +++ b/checker/hostname_test.go @@ -2,6 +2,7 @@ package checker import ( "bufio" + "context" "crypto/rand" "crypto/tls" "crypto/x509" @@ -100,7 +101,7 @@ func TestPolicyMatch(t *testing.T) { } func TestNoConnection(t *testing.T) { - result := FullCheckHostname("", "example.com", testTimeout) + result := FullCheckHostname(context.Background(), "", "example.com", testTimeout) expected := Result{ Status: 3, @@ -115,7 +116,7 @@ func TestNoTLS(t *testing.T) { ln := smtpListenAndServe(t, &tls.Config{}) defer ln.Close() - result := FullCheckHostname("", ln.Addr().String(), testTimeout) + result := FullCheckHostname(context.Background(), "", ln.Addr().String(), testTimeout) expected := Result{ Status: 2, @@ -135,7 +136,7 @@ func TestSelfSigned(t *testing.T) { ln := smtpListenAndServe(t, &tls.Config{Certificates: []tls.Certificate{cert}}) defer ln.Close() - result := FullCheckHostname("", ln.Addr().String(), testTimeout) + result := FullCheckHostname(context.Background(), "", ln.Addr().String(), testTimeout) expected := Result{ Status: 2, @@ -161,7 +162,7 @@ func TestNoTLS12(t *testing.T) { }) defer ln.Close() - result := FullCheckHostname("", ln.Addr().String(), testTimeout) + result := FullCheckHostname(context.Background(), "", ln.Addr().String(), testTimeout) expected := Result{ Status: 2, @@ -194,7 +195,7 @@ func TestSuccessWithFakeCA(t *testing.T) { // conserving the port number. addrParts := strings.Split(ln.Addr().String(), ":") port := addrParts[len(addrParts)-1] - result := FullCheckHostname("", "localhost:"+port, testTimeout) + result := FullCheckHostname(context.Background(), "", "localhost:"+port, testTimeout) expected := Result{ Status: 0, Checks: map[string]*Result{ @@ -217,7 +218,7 @@ func TestSuccessWithDelayedGreeting(t *testing.T) { defer ln.Close() go ServeDelayedGreeting(ln, t) - client, err := smtpDialWithTimeout(ln.Addr().String(), testTimeout) + client, err := smtpDialWithTimeout(context.Background(), ln.Addr().String(), testTimeout) if err != nil { t.Fatal(err) } @@ -269,7 +270,7 @@ func TestFailureWithBadHostname(t *testing.T) { // conserving the port number. addrParts := strings.Split(ln.Addr().String(), ":") port := addrParts[len(addrParts)-1] - result := FullCheckHostname("", "localhost:"+port, testTimeout) + result := FullCheckHostname(context.Background(), "", "localhost:"+port, testTimeout) expected := Result{ Status: 2, Checks: map[string]*Result{ @@ -309,7 +310,7 @@ func TestAdvertisedCiphers(t *testing.T) { ln := smtpListenAndServe(t, tlsConfig) defer ln.Close() - FullCheckHostname("", ln.Addr().String(), testTimeout) + FullCheckHostname(context.Background(), "", ln.Addr().String(), testTimeout) // Partial list of ciphers we want to support expectedCipherSuites := []struct { @@ -340,14 +341,13 @@ func compareStatuses(t *testing.T, expected Result, result HostnameResult) { if result.Status != expected.Status { t.Errorf("hostname status = %d, want %d", result.Status, expected.Status) } - if len(result.Checks) > len(expected.Checks) { t.Errorf("result contains too many checks\n expected %v\n want %v", result.Checks, expected.Checks) } - for _, c := range expected.Checks { - if got := result.Checks[c.Name].Status; got != c.Status { - t.Errorf("%s status = %d, want %d", c.Name, got, c.Status) + got, ok := result.Checks[c.Name] + if !ok || got.Status != c.Status { + t.Errorf("%s check result was %v, want status %d", c.Name, got, c.Status) } } } diff --git a/checker/mta_sts.go b/checker/mta_sts.go index cfc0f08b..e6895301 100644 --- a/checker/mta_sts.go +++ b/checker/mta_sts.go @@ -76,9 +76,9 @@ func getKeyValuePairs(record string, lineDelimiter string, return parsed } -func checkMTASTSRecord(domain string, timeout time.Duration) *Result { +func checkMTASTSRecord(ctx context.Context, domain string, timeout time.Duration) *Result { result := MakeResult(MTASTSText) - ctx, cancel := context.WithTimeout(context.Background(), timeout) + ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() var r net.Resolver records, err := r.LookupTXT(ctx, fmt.Sprintf("_mta-sts.%s", domain)) @@ -102,7 +102,7 @@ func validateMTASTSRecord(records []string, result *Result) *Result { return result.Success() } -func checkMTASTSPolicyFile(domain string, hostnameResults map[string]HostnameResult, timeout time.Duration) (*Result, string, map[string]string) { +func checkMTASTSPolicyFile(ctx context.Context, domain string, hostnameResults map[string]HostnameResult, timeout time.Duration) (*Result, string, map[string]string) { result := MakeResult(MTASTSPolicyFile) client := &http.Client{ Timeout: timeout, @@ -112,7 +112,10 @@ func checkMTASTSPolicyFile(domain string, hostnameResults map[string]HostnameRes }, } policyURL := fmt.Sprintf("https://mta-sts.%s/.well-known/mta-sts.txt", domain) - resp, err := client.Get(policyURL) + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + request, err := http.NewRequest("GET", policyURL, nil) + resp, err := client.Do(request.WithContext(ctx)) if err != nil { return result.Failure("Couldn't find policy file at %s.", policyURL), "", map[string]string{} } @@ -184,14 +187,14 @@ func validateMTASTSMXs(policyFileMXs []string, dnsMXs map[string]HostnameResult, } } -func (c Checker) checkMTASTS(domain string, hostnameResults map[string]HostnameResult) *MTASTSResult { +func (c Checker) checkMTASTS(ctx context.Context, domain string, hostnameResults map[string]HostnameResult) *MTASTSResult { if c.checkMTASTSOverride != nil { // Allow the Checker to mock this function. return c.checkMTASTSOverride(domain, hostnameResults) } result := MakeMTASTSResult() - result.addCheck(checkMTASTSRecord(domain, c.timeout())) - policyResult, policy, policyMap := checkMTASTSPolicyFile(domain, hostnameResults, c.timeout()) + result.addCheck(checkMTASTSRecord(ctx, domain, c.timeout())) + policyResult, policy, policyMap := checkMTASTSPolicyFile(ctx, domain, hostnameResults, c.timeout()) result.addCheck(policyResult) result.Policy = policy result.Mode = policyMap["mode"] diff --git a/checker/totals.go b/checker/totals.go index dd2049b2..c2e854c8 100644 --- a/checker/totals.go +++ b/checker/totals.go @@ -1,6 +1,7 @@ package checker import ( + "context" "encoding/csv" "io" "log" @@ -109,7 +110,7 @@ func (c *Checker) CheckCSV(domains *csv.Reader, resultHandler ResultHandler, dom for i := 0; i < poolSize; i++ { go func() { for domain := range work { - results <- c.CheckDomain(domain, nil) + results <- c.CheckDomain(context.TODO(), domain, nil) } done <- struct{}{} }() diff --git a/main.go b/main.go index e38b5f76..5650ecf3 100644 --- a/main.go +++ b/main.go @@ -38,7 +38,7 @@ func registerHandlers(api *API, mux *http.ServeMux) http.Handler { } // ServePublicEndpoints serves all public HTTP endpoints. -func ServePublicEndpoints(api *API, cfg *db.Config) { +func ServePublicEndpoints(ctx context.Context, exitSignal chan struct{}, api *API, cfg *db.Config) { mux := http.NewServeMux() mainHandler := registerHandlers(api, mux) @@ -52,20 +52,15 @@ func ServePublicEndpoints(api *API, cfg *db.Config) { Handler: mainHandler, } - exited := make(chan struct{}) go func() { - sigint := make(chan os.Signal, 1) - signal.Notify(sigint, os.Interrupt) - <-sigint - + <-ctx.Done() if err := server.Shutdown(context.Background()); err != nil { log.Printf("HTTP server Shutdown: %v", err) } - close(exited) }() log.Fatal(server.ListenAndServe()) - <-exited + exitSignal <- struct{}{} } // Loads a map of domains (effectively a set for fast lookup) to blacklist. @@ -114,14 +109,37 @@ func main() { Emailer: emailConfig, } api.parseTemplates() + + // Start go routines + exitSignals := make(chan struct{}) + ctx, cancel := context.WithCancel(context.Background()) + numRoutines := 0 if os.Getenv("VALIDATE_LIST") == "1" { log.Println("[Starting list validator]") - go validator.ValidateRegularly("Live policy list", list, 24*time.Hour) + go validator.ValidateRegularly(ctx, exitSignals, "Live policy list", list, 24*time.Hour) + numRoutines++ } if os.Getenv("VALIDATE_QUEUED") == "1" { log.Println("[Starting queued validator]") - go validator.ValidateRegularly("Testing domains", db, 24*time.Hour) + go validator.ValidateRegularly(ctx, exitSignals, "Testing domains", db, 24*time.Hour) + numRoutines++ + } + + go stats.UpdateRegularly(ctx, exitSignals, db, time.Hour) + go ServePublicEndpoints(ctx, exitSignals, &api, &cfg) + numRoutines += 2 + + // Register Ctrl+C signal handlers + sigint := make(chan os.Signal, 1) + signal.Notify(sigint, os.Interrupt) + go func() { + select { + case <-sigint: + cancel() + return + } + }() + for i := 0; i < numRoutines; i++ { + <-exitSignals } - go stats.UpdateRegularly(db, time.Hour) - ServePublicEndpoints(&api, &cfg) } diff --git a/main_test.go b/main_test.go index f470ff77..4b3df312 100644 --- a/main_test.go +++ b/main_test.go @@ -1,6 +1,7 @@ package main import ( + "context" "log" "net/http" "net/http/httptest" @@ -20,8 +21,8 @@ import ( var api *API var server *httptest.Server -func mockCheckPerform(message string) func(API, string) (checker.DomainResult, error) { - return func(api API, domain string) (checker.DomainResult, error) { +func mockCheckPerform(message string) func(context.Context, API, string) (checker.DomainResult, error) { + return func(ctx context.Context, api API, domain string) (checker.DomainResult, error) { return checker.NewSampleDomainResult(domain), nil } } diff --git a/scan_test.go b/scan_test.go index 02d4f61b..6a06050c 100644 --- a/scan_test.go +++ b/scan_test.go @@ -1,6 +1,7 @@ package main import ( + "context" "encoding/json" "io/ioutil" "net/http" @@ -126,7 +127,7 @@ func TestScanCached(t *testing.T) { data := url.Values{} data.Set("domain", "eff.org") http.PostForm(server.URL+"/api/scan", data) - original, _ := api.CheckDomain(*api, "eff.org") + original, _ := api.CheckDomain(context.Background(), *api, "eff.org") // Perform scan again, with different expected result. api.CheckDomain = mockCheckPerform("somethingelse") resp, _ := http.PostForm(server.URL+"/api/scan", data) diff --git a/stats/stats.go b/stats/stats.go index e0645ffa..e5150e7f 100644 --- a/stats/stats.go +++ b/stats/stats.go @@ -2,6 +2,7 @@ package stats import ( "bufio" + "context" "encoding/json" "fmt" "log" @@ -22,7 +23,7 @@ type Store interface { // Import imports aggregated scans from a remote server to the datastore. // Expected format is JSONL (newline-separated JSON objects). -func Import(store Store) error { +func Import(ctx context.Context, store Store) error { statsURL := os.Getenv("REMOTE_STATS_URL") resp, err := http.Get(statsURL) if err != nil { @@ -51,8 +52,8 @@ func Import(store Store) error { // Update imports aggregated scans and updates our cache table of local scans. // Log any errors. -func Update(store Store) { - err := Import(store) +func Update(ctx context.Context, store Store) { + err := Import(ctx, store) if err != nil { err = fmt.Errorf("Failed to import top domains stats: %v", err) log.Println(err) @@ -69,10 +70,17 @@ func Update(store Store) { } // UpdateRegularly runs Import to import aggregated stats from a remote server at regular intervals. -func UpdateRegularly(store Store, interval time.Duration) { +func UpdateRegularly(ctx context.Context, exited chan struct{}, store Store, interval time.Duration) { + ticker := time.NewTicker(interval) for { - Update(store) - <-time.After(interval) + select { + case <-ticker.C: + Update(ctx, store) + case <-ctx.Done(): + log.Printf("Shutting down stats updater...") + exited <- struct{}{} + return + } } } diff --git a/stats/stats_test.go b/stats/stats_test.go index 79ff577c..fc0ed7f4 100644 --- a/stats/stats_test.go +++ b/stats/stats_test.go @@ -1,6 +1,7 @@ package stats import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -58,7 +59,7 @@ func TestImport(t *testing.T) { defer ts.Close() os.Setenv("REMOTE_STATS_URL", ts.URL) store := mockAgScanStore{} - err := Import(&store) + err := Import(context.Background(), &store) if err != nil { t.Fatal(err) } @@ -79,7 +80,7 @@ func TestImport(t *testing.T) { func TestUpdate(t *testing.T) { store := mockAgScanStore{} - Update(&store) + Update(context.Background(), &store) a := store[0] // Confirm that date is trucated correctly if a.Time.Hour() != 0 || a.Time.Minute() != 0 { diff --git a/validator/validator.go b/validator/validator.go index 87d0b606..a94e991e 100644 --- a/validator/validator.go +++ b/validator/validator.go @@ -1,6 +1,7 @@ package validator import ( + "context" "fmt" "log" "time" @@ -28,7 +29,7 @@ func reportToSentry(name string, domain string, result checker.DomainResult) { result) } -type checkPerformer func(string, []string) checker.DomainResult +type checkPerformer func(context.Context, string, []string) checker.DomainResult type resultCallback func(string, string, checker.DomainResult) // Validator runs checks regularly against domain policies. This structure @@ -51,14 +52,14 @@ type Validator struct { checkPerformer checkPerformer } -func (v *Validator) checkPolicy(domain string, hostnames []string) checker.DomainResult { +func (v *Validator) checkPolicy(ctx context.Context, domain string, hostnames []string) checker.DomainResult { if v.checkPerformer == nil { c := checker.Checker{ Cache: checker.MakeSimpleCache(time.Hour), } v.checkPerformer = c.CheckDomain } - return v.checkPerformer(domain, hostnames) + return v.checkPerformer(ctx, domain, hostnames) } func (v *Validator) interval() time.Duration { @@ -81,30 +82,41 @@ func (v *Validator) policyPassed(name string, domain string, result checker.Doma } } -// Run starts the endless loop of validations. The first validation happens after the given -// Interval. Validation failures induce `policyFailed`, and successes cause `policyPassed`. -func (v *Validator) Run() { - for { - <-time.After(v.interval()) - log.Printf("[%s validator] starting regular validation", v.Name) - domains, err := v.Store.DomainsToValidate() +// Run performs a single validation. +// Validation failures induce `policyFailed`, and successes cause `policyPassed`. +func (v *Validator) runOnce(ctx context.Context) { + log.Printf("[%s validator] starting regular validation", v.Name) + domains, err := v.Store.DomainsToValidate() + if err != nil { + log.Printf("[%s validator] Could not retrieve domains: %v", v.Name, err) + return + } + for _, domain := range domains { + hostnames, err := v.Store.HostnamesForDomain(domain) if err != nil { - log.Printf("[%s validator] Could not retrieve domains: %v", v.Name, err) - continue + log.Printf("[%s validator] Could not retrieve policy for domain %s: %v", v.Name, domain, err) + return } - for _, domain := range domains { - hostnames, err := v.Store.HostnamesForDomain(domain) - if err != nil { - log.Printf("[%s validator] Could not retrieve policy for domain %s: %v", v.Name, domain, err) - continue - } - result := v.checkPolicy(domain, hostnames) - if result.Status != 0 { - log.Printf("[%s validator] %s failed; sending report", v.Name, domain) - v.policyFailed(v.Name, domain, result) - } else { - v.policyPassed(v.Name, domain, result) - } + result := v.checkPolicy(ctx, domain, hostnames) + if result.Status != 0 { + log.Printf("[%s validator] %s failed; sending report", v.Name, domain) + v.policyFailed(v.Name, domain, result) + } else { + v.policyPassed(v.Name, domain, result) + } + } +} + +func (v *Validator) runLoop(ctx context.Context, exited chan struct{}) { + ticker := time.NewTicker(v.Interval) + for { + select { + case <-ticker.C: + v.runOnce(ctx) + case <-ctx.Done(): + log.Printf("Shutting down %s validator...", v.Name) + exited <- struct{}{} + return } } } @@ -112,11 +124,11 @@ func (v *Validator) Run() { // ValidateRegularly regularly runs checker.CheckDomain against a Domain- // Hostname map. Interval specifies the interval to wait between each run. // Failures are reported to Sentry. -func ValidateRegularly(name string, store DomainPolicyStore, interval time.Duration) { +func ValidateRegularly(ctx context.Context, exited chan struct{}, name string, store DomainPolicyStore, interval time.Duration) { v := Validator{ Name: name, Store: store, Interval: interval, } - v.Run() + v.runLoop(ctx, exited) } diff --git a/validator/validator_test.go b/validator/validator_test.go index eec55a9f..88152308 100644 --- a/validator/validator_test.go +++ b/validator/validator_test.go @@ -1,6 +1,7 @@ package validator import ( + "context" "testing" "time" @@ -27,26 +28,30 @@ func noop(_ string, _ string, _ checker.DomainResult) {} func TestRegularValidationValidates(t *testing.T) { called := make(chan bool) - fakeChecker := func(domain string, hostnames []string) checker.DomainResult { + fakeChecker := func(_ context.Context, _ string, _ []string) checker.DomainResult { called <- true return checker.DomainResult{} } mock := mockDomainPolicyStore{ hostnames: map[string][]string{"a": []string{"hostname"}}} v := Validator{Store: mock, Interval: 100 * time.Millisecond, checkPerformer: fakeChecker, OnFailure: noop} - go v.Run() + ctx, cancel := context.WithCancel(context.Background()) + exited := make(chan struct{}) + go v.runLoop(ctx, exited) select { case <-called: - return + break case <-time.After(time.Second): t.Errorf("Checker wasn't called on hostname!") } + cancel() + <-exited } func TestRegularValidationReportsErrors(t *testing.T) { reports := make(chan string) - fakeChecker := func(domain string, hostnames []string) checker.DomainResult { + fakeChecker := func(_ context.Context, domain string, _ []string) checker.DomainResult { if domain == "fail" || domain == "error" { return checker.DomainResult{Status: 5} } @@ -67,15 +72,18 @@ func TestRegularValidationReportsErrors(t *testing.T) { v := Validator{Store: mock, Interval: 100 * time.Millisecond, checkPerformer: fakeChecker, OnFailure: fakeReporter, OnSuccess: fakeSuccessReporter, } - go v.Run() + ctx, cancel := context.WithCancel(context.Background()) + exited := make(chan struct{}) + go v.runLoop(ctx, exited) recvd := make(map[string]bool) numRecvd := 0 - for numRecvd < 4 { + for numRecvd < 6 { select { case report := <-successReports: if report != "normal" { t.Errorf("Didn't expect %s to succeed", report) } + numRecvd++ case report := <-reports: recvd[report] = true numRecvd++ @@ -83,6 +91,8 @@ func TestRegularValidationReportsErrors(t *testing.T) { t.Errorf("Timed out waiting for reports") } } + cancel() + <-exited if _, ok := recvd["fail"]; !ok { t.Errorf("Expected fail to be reported") }