tsweb: fix JSONHandler nil response

Signed-off-by: Zijie Lu <[email protected]>

tsweb/jsonhandler.go

@@ -113,7 +113,9 @@ func JSONHandler(fn interface{}) http.Handler {
 			}
 		case 2:
 			if vs[1].IsNil() {
-				writeResponse(w, http.StatusOK, responseSuccess(vs[0].Interface()))
+				if !vs[0].IsNil() {
+					writeResponse(w, http.StatusOK, responseSuccess(vs[0].Interface()))
+				}
 			} else {
 				err := vs[1].Interface().(error)
 				writeResponse(w, http.StatusBadRequest, responseError(err.Error()))

tsweb/jsonhandler_test.go

@@ -5,6 +5,7 @@
 package tsweb
 
 import (
+	"bytes"
 	"encoding/json"
 	"errors"
 	"net/http"
@@ -172,4 +173,24 @@ func TestNewJSONHandler(t *testing.T) {
 			return nil, "panic"
 		})
 	})
+
+	t.Run("2 2 forbidden", func(t *testing.T) {
+		code := http.StatusForbidden
+		body := []byte("forbidden")
+		h := JSONHandler(func(w http.ResponseWriter, r *http.Request) (*Data, error) {
+			w.WriteHeader(code)
+			w.Write(body)
+			return nil, nil
+		})
+
+		w := httptest.NewRecorder()
+		r := httptest.NewRequest("GET", "/", nil)
+		h.ServeHTTP(w, r)
+		if w.Code != http.StatusForbidden {
+			t.Fatalf("wrong code: %d %d", w.Code, code)
+		}
+		if !bytes.Equal(w.Body.Bytes(), []byte("forbidden")) {
+			t.Fatalf("wrong body: %s %s", w.Body.Bytes(), body)
+		}
+	})
 }