-
Notifications
You must be signed in to change notification settings - Fork 174
Expand file tree
/
Copy pathrest.go
More file actions
64 lines (56 loc) · 2.26 KB
/
rest.go
File metadata and controls
64 lines (56 loc) · 2.26 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
package httputil
import (
"context"
"encoding/json"
"fmt"
"net/http"
"github.com/stackrox/rox/pkg/logging"
"github.com/stackrox/rox/pkg/protocompat"
"google.golang.org/protobuf/encoding/protojson"
)
var (
log = logging.LoggerForModule()
)
type responseHeaderContextKey struct{}
// RESTHandler wraps a function implementing a REST endpoint as an http.Handler. The function receives the incoming HTTP
// request, and returns a response object or an error. In case of a non-nil error, this error is written in JSON/gRPC
// style via `WriteError` (i.e., you can return a plain Golang error (equivalent to status code 500), a gRPC-style
// status, or an `HTTPError`). Otherwise, the returned object is written as JSON (using `jsonpb` if it is a
// `protocompat.Message`, and `encoding/json` otherwise).
// If you need to mutate response headers, these can be accessed by calling ResponseHeaderFromContext.
// Any errors that occur writing to the response body are simply logged.
func RESTHandler(endpointFunc func(*http.Request) (interface{}, error)) http.HandlerFunc {
return func(w http.ResponseWriter, req *http.Request) {
w.Header().Set("Content-Type", "application/json")
reqCtx := context.WithValue(req.Context(), responseHeaderContextKey{}, w.Header())
resp, err := endpointFunc(req.WithContext(reqCtx))
if err != nil {
WriteError(w, err)
return
}
if resp == nil {
_, err = fmt.Fprint(w, "{}")
} else if protoMsg, _ := resp.(protocompat.Message); protoMsg != nil {
var data []byte
data, err = protojson.Marshal(protoMsg)
if err != nil {
log.Errorf("Failed to send response from REST handler: %v", err)
}
_, err = w.Write(data)
if err != nil {
log.Errorf("Failed to send response from REST handler: %v", err)
}
} else {
err = json.NewEncoder(w).Encode(resp)
}
if err != nil {
log.Errorf("Failed to send response from REST handler: %v", err)
}
}
}
// ResponseHeaderFromContext returns the (mutable) response header from a given context. This only works from within a
// REST endpoint function wrapped with RESTHandler, but is guaranteed to always return a non-nil result in this case.
func ResponseHeaderFromContext(ctx context.Context) http.Header {
hdr, _ := ctx.Value(responseHeaderContextKey{}).(http.Header)
return hdr
}