diff --git a/api/api.go b/api/api.go index 8dd4973..b8e790e 100644 --- a/api/api.go +++ b/api/api.go @@ -19,8 +19,10 @@ import ( ) const ( - IP_HEADER = "x-ifconfig-ip" - COUNTRY_HEADER = "x-ifconfig-country" + IP_HEADER = "x-ifconfig-ip" + COUNTRY_HEADER = "x-ifconfig-country" + TEXT_PLAIN = "text/plain; charset=utf-8" + APPLICATION_JSON = "application/json" ) var cliUserAgentExp = regexp.MustCompile("^(?i)(curl|wget|fetch\\slibfetch)\\/.*$") @@ -29,7 +31,6 @@ type API struct { db *geoip2.Reader CORS bool Template string - Logger *log.Logger } func New() *API { return &API{} } @@ -84,13 +85,17 @@ func ipFromRequest(r *http.Request) (net.IP, error) { return ip, nil } -func headerKeyFromRequest(r *http.Request) string { +func headerPairFromRequest(r *http.Request) (string, string, error) { vars := mux.Vars(r) - key, ok := vars["key"] + header, ok := vars["header"] if !ok { - return "" + header = IP_HEADER } - return key + value := r.Header.Get(header) + if value == "" { + return "", "", fmt.Errorf("no value found for: %s", header) + } + return header, value, nil } func (a *API) lookupCountry(ip net.IP) (string, error) { @@ -110,29 +115,16 @@ func (a *API) lookupCountry(ip net.IP) (string, error) { return "", fmt.Errorf("could not determine country for IP: %s", ip) } -func (a *API) handleError(w http.ResponseWriter, err error) { - a.Logger.Print(err) - w.WriteHeader(http.StatusInternalServerError) - io.WriteString(w, "Internal server error") -} - -func handleUnknownHeader(w http.ResponseWriter, key string) { - w.WriteHeader(http.StatusBadRequest) - io.WriteString(w, "Bad request: Unknown header: "+key) -} - -func (a *API) DefaultHandler(w http.ResponseWriter, r *http.Request) { +func (a *API) DefaultHandler(w http.ResponseWriter, r *http.Request) *appError { cmd := cmdFromQueryParams(r.URL.Query()) funcMap := template.FuncMap{"ToLower": strings.ToLower} t, err := template.New(filepath.Base(a.Template)).Funcs(funcMap).ParseFiles(a.Template) if err != nil { - a.handleError(w, err) - return + return internalServerError(err) } b, err := json.MarshalIndent(r.Header, "", " ") if err != nil { - a.handleError(w, err) - return + return internalServerError(err) } var data = struct { @@ -143,42 +135,37 @@ func (a *API) DefaultHandler(w http.ResponseWriter, r *http.Request) { }{r.Header.Get(IP_HEADER), string(b), r.Header, cmd} if err := t.Execute(w, &data); err != nil { - a.handleError(w, err) + return internalServerError(err) } + return nil } -func (a *API) JSONHandler(w http.ResponseWriter, r *http.Request) { - key := headerKeyFromRequest(r) - if key == "" { - key = IP_HEADER - } - value := map[string]string{key: r.Header.Get(key)} - if value[key] == "" { - handleUnknownHeader(w, key) - return +func (a *API) JSONHandler(w http.ResponseWriter, r *http.Request) *appError { + k, v, err := headerPairFromRequest(r) + contentType := APPLICATION_JSON + if err != nil { + return notFound(err).WithContentType(contentType) } + value := map[string]string{k: v} b, err := json.MarshalIndent(value, "", " ") if err != nil { - a.handleError(w, err) - return + return internalServerError(err).WithContentType(contentType) } + w.Header().Set("Content-Type", contentType) w.Write(b) + return nil } -func (a *API) CLIHandler(w http.ResponseWriter, r *http.Request) { - key := headerKeyFromRequest(r) - if key == "" { - key = IP_HEADER +func (a *API) CLIHandler(w http.ResponseWriter, r *http.Request) *appError { + _, v, err := headerPairFromRequest(r) + if err != nil { + return notFound(err) } - value := r.Header.Get(key) - if value == "" { - handleUnknownHeader(w, key) - return + if !strings.HasSuffix(v, "\n") { + v += "\n" } - if !strings.HasSuffix(value, "\n") { - value += "\n" - } - io.WriteString(w, value) + io.WriteString(w, v) + return nil } func cliMatcher(r *http.Request, rm *mux.RouteMatch) bool { @@ -207,20 +194,51 @@ func (a *API) requestFilter(next http.Handler) http.Handler { }) } +type appHandler func(http.ResponseWriter, *http.Request) *appError + +func (fn appHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if e := fn(w, r); e != nil { // e is *appError + if e.Error != nil { + log.Print(e.Error) + } + contentType := e.ContentType + if contentType == "" { + contentType = TEXT_PLAIN + } + response := e.Response + if response == "" { + response = e.Error.Error() + } + if e.IsJSON() { + var data = struct { + Error string `json:"error"` + }{response} + b, err := json.MarshalIndent(data, "", " ") + if err != nil { + panic(err) + } + response = string(b) + } + w.Header().Set("Content-Type", contentType) + w.WriteHeader(e.Code) + io.WriteString(w, response) + } +} + func (a *API) Handlers() http.Handler { r := mux.NewRouter() // JSON - r.HandleFunc("/", a.JSONHandler).Methods("GET").Headers("Accept", "application/json") - r.HandleFunc("/{key}", a.JSONHandler).Methods("GET").Headers("Accept", "application/json") - r.HandleFunc("/{key}.json", a.JSONHandler).Methods("GET") + r.Handle("/", appHandler(a.JSONHandler)).Methods("GET").Headers("Accept", APPLICATION_JSON) + r.Handle("/{header}", appHandler(a.JSONHandler)).Methods("GET").Headers("Accept", APPLICATION_JSON) + r.Handle("/{header}.json", appHandler(a.JSONHandler)).Methods("GET") // CLI - r.HandleFunc("/", a.CLIHandler).Methods("GET").MatcherFunc(cliMatcher) - r.HandleFunc("/{key}", a.CLIHandler).Methods("GET").MatcherFunc(cliMatcher) + r.Handle("/", appHandler(a.CLIHandler)).Methods("GET").MatcherFunc(cliMatcher) + r.Handle("/{header}", appHandler(a.CLIHandler)).Methods("GET").MatcherFunc(cliMatcher) // Default - r.HandleFunc("/", a.DefaultHandler).Methods("GET") + r.Handle("/", appHandler(a.DefaultHandler)).Methods("GET") // Pass all requests through the request filter return a.requestFilter(r) diff --git a/api/api_test.go b/api/api_test.go index 01fe0e2..94ca899 100644 --- a/api/api_test.go +++ b/api/api_test.go @@ -3,6 +3,7 @@ package api import ( "fmt" "io/ioutil" + "log" "net" "net/http" "net/http/httptest" @@ -11,10 +12,10 @@ import ( "testing" ) -func httpGet(url string, json bool, userAgent string) (string, error) { +func httpGet(url string, json bool, userAgent string) (string, int, error) { r, err := http.NewRequest("GET", url, nil) if err != nil { - return "", err + return "", 0, err } if json { r.Header.Set("Accept", "application/json") @@ -22,17 +23,18 @@ func httpGet(url string, json bool, userAgent string) (string, error) { r.Header.Set("User-Agent", userAgent) res, err := http.DefaultClient.Do(r) if err != nil { - return "", err + return "", 0, err } defer res.Body.Close() data, err := ioutil.ReadAll(res.Body) if err != nil { - return "", err + return "", 0, err } - return string(data), nil + return string(data), res.StatusCode, nil } func TestGetIP(t *testing.T) { + log.SetOutput(ioutil.Discard) toJSON := func(k string, v string) string { return fmt.Sprintf("{\n \"%s\": \"%s\"\n}", k, v) } @@ -42,18 +44,24 @@ func TestGetIP(t *testing.T) { json bool out string userAgent string + status int }{ - {s.URL, false, "127.0.0.1\n", "curl/7.26.0"}, - {s.URL, false, "127.0.0.1\n", "Wget/1.13.4 (linux-gnu)"}, - {s.URL, false, "127.0.0.1\n", "fetch libfetch/2.0"}, - {s.URL + "/x-ifconfig-ip.json", false, toJSON("x-ifconfig-ip", "127.0.0.1"), ""}, - {s.URL, true, toJSON("x-ifconfig-ip", "127.0.0.1"), ""}, + {s.URL, false, "127.0.0.1\n", "curl/7.26.0", 200}, + {s.URL, false, "127.0.0.1\n", "Wget/1.13.4 (linux-gnu)", 200}, + {s.URL, false, "127.0.0.1\n", "fetch libfetch/2.0", 200}, + {s.URL + "/x-ifconfig-ip.json", false, toJSON("x-ifconfig-ip", "127.0.0.1"), "", 200}, + {s.URL, true, toJSON("x-ifconfig-ip", "127.0.0.1"), "", 200}, + {s.URL + "/foo", false, "no value found for: foo", "curl/7.26.0", 404}, + {s.URL + "/foo", true, "{\n \"error\": \"no value found for: foo\"\n}", "curl/7.26.0", 404}, } for _, tt := range tests { - out, err := httpGet(tt.url, tt.json, tt.userAgent) + out, status, err := httpGet(tt.url, tt.json, tt.userAgent) if err != nil { t.Fatal(err) } + if status != tt.status { + t.Errorf("Expected %d, got %d", tt.status, status) + } if out != tt.out { t.Errorf("Expected %q, got %q", tt.out, out) } diff --git a/api/error.go b/api/error.go new file mode 100644 index 0000000..e688c8d --- /dev/null +++ b/api/error.go @@ -0,0 +1,42 @@ +package api + +import "net/http" + +type appError struct { + Error error + Response string + Code int + ContentType string +} + +func internalServerError(err error) *appError { + return &appError{Error: err, Response: "Internal server error", Code: http.StatusInternalServerError} +} + +func notFound(err error) *appError { + return &appError{Error: err, Code: http.StatusNotFound} +} + +func (e *appError) WithContentType(contentType string) *appError { + e.ContentType = contentType + return e +} + +func (e *appError) WithCode(code int) *appError { + e.Code = code + return e +} + +func (e *appError) WithResponse(response string) *appError { + e.Response = response + return e +} + +func (e *appError) WithError(err error) *appError { + e.Error = err + return e +} + +func (e *appError) IsJSON() bool { + return e.ContentType == APPLICATION_JSON +} diff --git a/main.go b/main.go index cffc84a..83fef92 100644 --- a/main.go +++ b/main.go @@ -33,9 +33,8 @@ func main() { a.CORS = opts.CORS a.Template = opts.Template - a.Logger = log.New(os.Stderr, "", log.LstdFlags) - a.Logger.Printf("Listening on %s", opts.Listen) + log.Printf("Listening on %s", opts.Listen) if err := a.ListenAndServe(opts.Listen); err != nil { log.Fatal(err) }