rawlinp closed pull request #3171: Add TO api helper safety to prevent double
writes
URL: https://github.com/apache/trafficcontrol/pull/3171
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/traffic_ops/traffic_ops_golang/api/api.go
b/traffic_ops/traffic_ops_golang/api/api.go
index d1ec4f29f..7a4bb7c37 100644
--- a/traffic_ops/traffic_ops_golang/api/api.go
+++ b/traffic_ops/traffic_ops_golang/api/api.go
@@ -46,6 +46,7 @@ import (
const DBContextKey = "db"
const ConfigContextKey = "context"
const ReqIDContextKey = "reqid"
+const APIRespWrittenKey = "respwritten"
type CRUDFactory func(reqInfo *APIInfo) CRUDer
@@ -60,6 +61,12 @@ func WriteResp(w http.ResponseWriter, r *http.Request, v
interface{}) {
// WriteRespRaw acts like WriteResp, but doesn't wrap the object in a
`{"response":` object. This should be used to respond with endpoints which
don't wrap their response in a "response" object.
func WriteRespRaw(w http.ResponseWriter, r *http.Request, v interface{}) {
+ if respWritten(r) {
+ log.Errorf("WriteRespRaw called after a write already occurred!
Not double-writing! Path %s", r.URL.Path)
+ return
+ }
+ setRespWritten(r)
+
bts, err := json.Marshal(v)
if err != nil {
log.Errorf("marshalling JSON (raw) for %T: %v", v, err)
@@ -73,6 +80,12 @@ func WriteRespRaw(w http.ResponseWriter, r *http.Request, v
interface{}) {
// WriteRespVals is like WriteResp, but also takes a map of root-level values
to write. The API most commonly needs these for meta-parameters, like size,
limit, and orderby.
// This is a helper for the common case; not using this in unusual cases is
perfectly acceptable.
func WriteRespVals(w http.ResponseWriter, r *http.Request, v interface{}, vals
map[string]interface{}) {
+ if respWritten(r) {
+ log.Errorf("WriteRespVals called after a write already
occurred! Not double-writing! Path %s", r.URL.Path)
+ return
+ }
+ setRespWritten(r)
+
vals["response"] = v
respBts, err := json.Marshal(vals)
if err != nil {
@@ -90,6 +103,12 @@ func WriteRespVals(w http.ResponseWriter, r *http.Request,
v interface{}, vals m
//
// This is a helper for the common case; not using this in unusual cases is
perfectly acceptable.
func HandleErr(w http.ResponseWriter, r *http.Request, tx *sql.Tx, statusCode
int, userErr error, sysErr error) {
+ if respWritten(r) {
+ log.Errorf("HandleErr called after a write already occurred!
Attempting to write the error anyway! Path %s", r.URL.Path)
+ // Don't return, attempt to rollback and write the error anyway
+ }
+ setRespWritten(r)
+
if tx != nil {
if err := tx.Rollback(); err != nil && err != sql.ErrTxDone {
log.Errorln("rolling back transaction: " + err.Error())
@@ -146,6 +165,12 @@ func RespWriterVals(w http.ResponseWriter, r
*http.Request, tx *sql.Tx, vals map
// WriteRespAlert creates an alert, serializes it as JSON, and writes that to
w. Any errors are logged and written to w via tc.GetHandleErrorsFunc.
// This is a helper for the common case; not using this in unusual cases is
perfectly acceptable.
func WriteRespAlert(w http.ResponseWriter, r *http.Request, level
tc.AlertLevel, msg string) {
+ if respWritten(r) {
+ log.Errorf("WriteRespAlert called after a write already
occurred! Not double-writing! Path %s", r.URL.Path)
+ return
+ }
+ setRespWritten(r)
+
resp := struct{ tc.Alerts }{tc.CreateAlerts(level, msg)}
respBts, err := json.Marshal(resp)
if err != nil {
@@ -159,6 +184,12 @@ func WriteRespAlert(w http.ResponseWriter, r
*http.Request, level tc.AlertLevel,
// WriteRespAlertObj Writes the given alert, and the given response object.
// This is a helper for the common case; not using this in unusual cases is
perfectly acceptable.
func WriteRespAlertObj(w http.ResponseWriter, r *http.Request, level
tc.AlertLevel, msg string, obj interface{}) {
+ if respWritten(r) {
+ log.Errorf("WriteRespAlertObj called after a write already
occurred! Not double-writing! Path %s", r.URL.Path)
+ return
+ }
+ setRespWritten(r)
+
resp := struct {
tc.Alerts
Response interface{} `json:"response"`
@@ -376,6 +407,19 @@ func getReqID(ctx context.Context) (uint64, error) {
return 0, errors.New("No ReqID found in Context")
}
+// setRespWritten sets the APIRespWrittenKey key in the Context of the given
Request.
+// This is used to indicate that a response has been written with an API
helper, and to prevent double-write errors.
+// If an API helper which responds is called after another response helper was
already called, all API helpers will log an error, and not write the second
response, except HandleErr, which will write its error anyway, along with its
status code.
+func setRespWritten(r *http.Request) {
+ *r = *r.WithContext(context.WithValue(r.Context(), APIRespWrittenKey,
struct{}{}))
+}
+
+// respWritten gets the APIRespWrittenKey key, which indicates whether an API
response helper was previously called.
+// This is used to prevent double-write errors. See setRespWritten.
+func respWritten(r *http.Request) bool {
+ return r.Context().Value(APIRespWrittenKey) != nil
+}
+
// TypeErrToAPIErr takes a slice of errors and an ApiErrorType, and converts
them to the (userErr, sysErr, errCode) idiom used by the api package.
func TypeErrsToAPIErr(errs []error, errType tc.ApiErrorType) (error, error,
int) {
if len(errs) == 0 {
diff --git a/traffic_ops/traffic_ops_golang/api/api_test.go
b/traffic_ops/traffic_ops_golang/api/api_test.go
index c9df510e2..7b65d8c61 100644
--- a/traffic_ops/traffic_ops_golang/api/api_test.go
+++ b/traffic_ops/traffic_ops_golang/api/api_test.go
@@ -19,10 +19,18 @@ package api
* under the License.
*/
-import "testing"
+import (
+ "database/sql"
+ "encoding/json"
+ "errors"
+ "net/http"
+ "net/url"
+ "testing"
-func TestCamelCase(t *testing.T) {
+ "github.com/apache/trafficcontrol/lib/go-tc"
+)
+func TestCamelCase(t *testing.T) {
testStrings := []string{"hello_world", "trailing_underscore_",
"w_h_a_t____"}
expected := []string{"helloWorld", "trailingUnderscore", "wHAT"}
for i, str := range testStrings {
@@ -31,3 +39,126 @@ func TestCamelCase(t *testing.T) {
}
}
}
+
+// TestRespWrittenAfterErrFails tests that a WriteResp called after HandleErr
will not be written.
+func TestRespWrittenAfterErrFails(t *testing.T) {
+ w := &MockHTTPResponseWriter{}
+ r := &http.Request{URL: &url.URL{}}
+ tx := (*sql.Tx)(nil)
+
+ expectedCode := http.StatusUnauthorized
+ expectedUserErr := errors.New("user unauthorized")
+
+ HandleErr(w, r, tx, expectedCode, expectedUserErr, nil)
+ WriteResp(w, r, "should not be written")
+
+ actualCode := w.Code
+ statusVal := r.Context().Value(tc.StatusKey)
+ statusInt, ok := statusVal.(int)
+ if ok {
+ actualCode = statusInt
+ }
+
+ if actualCode != expectedCode {
+ t.Errorf("code expected: %+v, actual %+v", expectedCode,
actualCode)
+ }
+
+ alerts := tc.Alerts{}
+ if err := json.Unmarshal(w.Body, &alerts); err != nil {
+ t.Fatalf("unmarshalling actual body: %v", err)
+ }
+ for _, alert := range alerts.Alerts {
+ if string(alert.Level) != tc.ErrorLevel.String() {
+ t.Errorf("alert level expected: '%s', actual: '%s'",
tc.ErrorLevel.String(), alert.Level)
+ }
+ }
+}
+
+func TestWriteResp(t *testing.T) {
+ apiWriteTest(t, func(w http.ResponseWriter, r *http.Request) {
+ WriteResp(w, r, "foo")
+ })
+}
+
+func TestWriteRespRaw(t *testing.T) {
+ apiWriteTest(t, func(w http.ResponseWriter, r *http.Request) {
+ WriteRespRaw(w, r, "foo")
+ })
+}
+
+func TestWriteRespVals(t *testing.T) {
+ apiWriteTest(t, func(w http.ResponseWriter, r *http.Request) {
+ WriteRespVals(w, r, "foo", map[string]interface{}{"a": "b"})
+ })
+}
+
+func TestRespWriter(t *testing.T) {
+ apiWriteTest(t, func(w http.ResponseWriter, r *http.Request) {
+ RespWriter(w, r, nil)("foo", nil)
+ })
+}
+
+func TestRespWriterVals(t *testing.T) {
+ apiWriteTest(t, func(w http.ResponseWriter, r *http.Request) {
+ RespWriterVals(w, r, nil, map[string]interface{}{"a":
"b"})("foo", nil)
+ })
+}
+
+func TestWriteRespAlert(t *testing.T) {
+ apiWriteTest(t, func(w http.ResponseWriter, r *http.Request) {
+ WriteRespAlert(w, r, tc.ErrorLevel, "foo error")
+ })
+}
+
+func TestWriteRespAlertObj(t *testing.T) {
+ apiWriteTest(t, func(w http.ResponseWriter, r *http.Request) {
+ WriteRespAlertObj(w, r, tc.ErrorLevel, "foo error", "bar")
+ })
+}
+
+// apiWriteTest tests that an API write func succeeds and writes a body and a
200.
+func apiWriteTest(t *testing.T, write func(w http.ResponseWriter, r
*http.Request)) {
+ w := &MockHTTPResponseWriter{}
+ r := &http.Request{URL: &url.URL{}}
+
+ write(w, r)
+
+ if w.Code == 0 {
+ w.Code = http.StatusOK // emulate behavior of w.Write
+ }
+
+ actualCode := w.Code
+ statusVal := r.Context().Value(tc.StatusKey)
+ statusInt, ok := statusVal.(int)
+ if ok {
+ actualCode = statusInt
+ }
+
+ expectedCode := http.StatusOK
+
+ if actualCode != expectedCode {
+ t.Errorf("code expected: %+v, actual %+v", expectedCode,
actualCode)
+ }
+
+ if len(w.Body) == 0 {
+ t.Errorf("body len expected: >0, actual 0")
+ }
+}
+
+type MockHTTPResponseWriter struct {
+ Code int
+ Body []byte
+}
+
+func (i *MockHTTPResponseWriter) WriteHeader(rc int) {
+ i.Code = rc
+}
+
+func (i *MockHTTPResponseWriter) Write(b []byte) (int, error) {
+ i.Body = append(i.Body, b...)
+ return len(b), nil
+}
+
+func (i *MockHTTPResponseWriter) Header() http.Header {
+ return http.Header{}
+}
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
