diff --git a/broker/broker.go b/broker/broker.go index 4db1511..361040a 100644 --- a/broker/broker.go +++ b/broker/broker.go @@ -10,6 +10,7 @@ import ( "crypto/tls" "flag" "fmt" + "io" "io/ioutil" "log" "net" @@ -37,10 +38,10 @@ type BrokerContext struct { metrics *Metrics } -func NewBrokerContext() *BrokerContext { +func NewBrokerContext(metricsLogger *log.Logger) *BrokerContext { snowflakes := new(SnowflakeHeap) heap.Init(snowflakes) - metrics, err := NewMetrics() + metrics, err := NewMetrics(metricsLogger) if err != nil { panic(err.Error()) @@ -253,6 +254,7 @@ func main() { var geoip6Database string var disableTLS bool var disableGeoip bool + var metricsFilename string flag.StringVar(&acmeEmail, "acme-email", "", "optional contact email for Let's Encrypt notifications") flag.StringVar(&acmeHostnamesCommas, "acme-hostnames", "", "comma-separated hostnames for TLS certificate") @@ -261,11 +263,27 @@ func main() { flag.StringVar(&geoip6Database, "geoip6db", "/usr/share/tor/geoip6", "path to correctly formatted geoip database mapping IPv6 address ranges to country codes") flag.BoolVar(&disableTLS, "disable-tls", false, "don't use HTTPS") flag.BoolVar(&disableGeoip, "disable-geoip", false, "don't use geoip for stats collection") + flag.StringVar(&metricsFilename, "metrics-log", "", "path to metrics logging output") flag.Parse() + var metricsFile io.Writer = os.Stdout + var err error + log.SetFlags(log.LstdFlags | log.LUTC) - ctx := NewBrokerContext() + if metricsFilename != "" { + metricsFile, err = os.OpenFile(metricsFilename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) + + if err != nil { + log.Fatal(err.Error()) + } + } else { + metricsFile = os.Stdout + } + + metricsLogger := log.New(metricsFile, "", log.LstdFlags|log.LUTC) + + ctx := NewBrokerContext(metricsLogger) if !disableGeoip { err := ctx.metrics.LoadGeoipDatabases(geoipDatabase, geoip6Database) @@ -283,7 +301,6 @@ func main() { http.Handle("/answer", SnowflakeHandler{ctx, proxyAnswers}) http.Handle("/debug", SnowflakeHandler{ctx, debugHandler}) - var err error server := http.Server{ Addr: addr, } diff --git a/broker/geoip.go b/broker/geoip.go index 27a87fd..16cada1 100644 --- a/broker/geoip.go +++ b/broker/geoip.go @@ -209,7 +209,8 @@ func GeoIPLoadFile(table GeoIPTable, pathname string) error { return nil } -//Returns the country location of an IPv4 or IPv6 address. +//Returns the country location of an IPv4 or IPv6 address, and a boolean value +//that indicates whether the IP address was present in the geoip database func GetCountryByAddr(table GeoIPTable, ip net.IP) (string, bool) { table.Lock() diff --git a/broker/metrics.go b/broker/metrics.go index 9eabe57..f961d1f 100644 --- a/broker/metrics.go +++ b/broker/metrics.go @@ -5,7 +5,6 @@ import ( "fmt" "log" "net" - "os" "sync" "time" ) @@ -28,7 +27,7 @@ type Metrics struct { } func (s CountryStats) Display() string { - return fmt.Sprintln(s.counts) + return fmt.Sprint(s.counts) } func (m *Metrics) UpdateCountryStats(addr string) { @@ -56,9 +55,7 @@ func (m *Metrics) UpdateCountryStats(addr string) { } //update map of countries and counts - if country != "" { - m.countryStats.counts[country]++ - } + m.countryStats.counts[country]++ return } @@ -88,17 +85,9 @@ func (m *Metrics) LoadGeoipDatabases(geoipDB string, geoip6DB string) error { return nil } -func NewMetrics() (*Metrics, error) { +func NewMetrics(metricsLogger *log.Logger) (*Metrics, error) { m := new(Metrics) - f, err := os.OpenFile("metrics.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - - if err != nil { - return nil, err - } - - metricsLogger := log.New(f, "", log.LstdFlags|log.LUTC) - m.countryStats = CountryStats{ counts: make(map[string]int), }