fix: fix error handler middleware

cmd/main/main.go

@@ -171,7 +171,7 @@ func main() {
 
 	gwS := runtime.NewServeMux(
 		runtime.WithForwardResponseOption(httpResponseModifier),
-		runtime.WithErrorHandler(handleError),
+		runtime.WithErrorHandler(errorHandler),
 		runtime.WithIncomingHeaderMatcher(customMatcher),
 		runtime.WithMarshalerOption(runtime.MIMEWildcard, &runtime.JSONPb{
 			MarshalOptions: protojson.MarshalOptions{

cmd/main/misc.go

@@ -2,15 +2,17 @@ package main
 
 import (
 	"context"
-	"encoding/json"
+	"fmt"
+	"io"
 	"net/http"
+	"net/textproto"
 	"strconv"
 	"strings"
 
-	"github.com/gogo/status"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
-	"github.com/instill-ai/connector-backend/pkg/datamodel"
+	"github.com/instill-ai/connector-backend/internal/logger"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -35,35 +37,92 @@ func httpResponseModifier(ctx context.Context, w http.ResponseWriter, p proto.Me
 	return nil
 }
 
-func handleError(ctx context.Context, mux *runtime.ServeMux, marshaler runtime.Marshaler, w http.ResponseWriter, r *http.Request, err error) {
-	if s, ok := status.FromError(err); ok {
-		switch {
-		case s.Code() == codes.FailedPrecondition && strings.Contains(s.Message(), "[DELETE]"):
-			errorResponse(w,
-				http.StatusUnprocessableEntity,
-				http.StatusText(http.StatusUnprocessableEntity),
-				s.Message(),
-			)
-		case s.Code() == codes.FailedPrecondition:
-			errorResponse(w,
-				http.StatusPreconditionFailed,
-				http.StatusText(http.StatusPreconditionFailed),
-				s.Message(),
-			)
-		case s.Code() == codes.InvalidArgument:
-			errorResponse(w,
-				http.StatusBadRequest,
-				http.StatusText(http.StatusBadRequest),
-				s.Message(),
-			)
-		case s.Code() == codes.AlreadyExists:
-			errorResponse(w,
-				http.StatusConflict,
-				http.StatusText(http.StatusConflict),
-				s.Message(),
-			)
+func errorHandler(ctx context.Context, mux *runtime.ServeMux, marshaler runtime.Marshaler, w http.ResponseWriter, r *http.Request, err error) {
+
+	logger, _ := logger.GetZapLogger()
+
+	// return Internal when Marshal failed
+	const fallback = `{"code": 13, "message": "failed to marshal error message"}`
+
+	s := status.Convert(err)
+	pb := s.Proto()
+
+	w.Header().Del("Trailer")
+	w.Header().Del("Transfer-Encoding")
+
+	contentType := marshaler.ContentType(pb)
+	if contentType == "application/json" {
+		w.Header().Set("Content-Type", "application/problem+json")
+	} else {
+		w.Header().Set("Content-Type", contentType)
+	}
+
+	if s.Code() == codes.Unauthenticated {
+		w.Header().Set("WWW-Authenticate", s.Message())
+	}
+
+	buf, err := marshaler.Marshal(pb)
+	if err != nil {
+		logger.Info(fmt.Sprintf("Failed to marshal error message %q: %v", s, err))
+		w.WriteHeader(http.StatusInternalServerError)
+		if _, err := io.WriteString(w, fallback); err != nil {
+			logger.Info(fmt.Sprintf("Failed to write response: %v", err))
+		}
+		return
+	}
+
+	md, ok := runtime.ServerMetadataFromContext(ctx)
+	if !ok {
+		logger.Info("Failed to extract ServerMetadata from context")
+	}
+
+	for k, vs := range md.HeaderMD {
+		if h, ok := func(key string) (string, bool) {
+			return fmt.Sprintf("%s%s", runtime.MetadataHeaderPrefix, key), true
+		}(k); ok {
+			for _, v := range vs {
+				w.Header().Add(h, v)
+			}
 		}
 	}
+
+	// RFC 7230 https://tools.ietf.org/html/rfc7230#section-4.1.2
+	// Unless the request includes a TE header field indicating "trailers"
+	// is acceptable, as described in Section 4.3, a server SHOULD NOT
+	// generate trailer fields that it believes are necessary for the user
+	// agent to receive.
+	doForwardTrailers := strings.Contains(strings.ToLower(r.Header.Get("TE")), "trailers")
+
+	if doForwardTrailers {
+		for k := range md.TrailerMD {
+			tKey := textproto.CanonicalMIMEHeaderKey(fmt.Sprintf("%s%s", runtime.MetadataTrailerPrefix, k))
+			w.Header().Add("Trailer", tKey)
+		}
+		w.Header().Set("Transfer-Encoding", "chunked")
+	}
+
+	var st int
+	switch {
+	case s.Code() == codes.FailedPrecondition && strings.Contains(s.Message(), "[DELETE]"):
+		st = http.StatusUnprocessableEntity
+	default:
+		st = runtime.HTTPStatusFromCode(s.Code())
+	}
+
+	w.WriteHeader(st)
+	if _, err := w.Write(buf); err != nil {
+		logger.Info(fmt.Sprintf("Failed to write response: %v", err))
+	}
+
+	if doForwardTrailers {
+		for k, vs := range md.TrailerMD {
+			tKey := fmt.Sprintf("%s%s", runtime.MetadataTrailerPrefix, k)
+			for _, v := range vs {
+				w.Header().Add(tKey, v)
+			}
+		}
+	}
+
 }
 
 func customMatcher(key string) (string, bool) {
@@ -78,14 +137,3 @@ func customMatcher(key string) (string, bool) {
 		return runtime.DefaultHeaderMatcher(key)
 	}
 }
-
-func errorResponse(w http.ResponseWriter, status int, title string, detail string) {
-	w.Header().Add("Content-Type", "application/json+problem")
-	w.WriteHeader(status)
-	obj, _ := json.Marshal(datamodel.Error{
-		Status: int32(status),
-		Title:  title,
-		Detail: detail,
-	})
-	_, _ = w.Write(obj)
-}

go.mod

@@ -40,7 +40,7 @@ require (
 	github.com/go-sql-driver/mysql v1.6.0 // indirect
 	github.com/gogo/googleapis v1.4.1 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
-	github.com/gogo/status v1.1.0
+	github.com/gogo/status v1.1.0 // indirect
 	github.com/golang/mock v1.6.0 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/google/uuid v1.3.0 // indirect

go.sum

@@ -509,21 +509,6 @@ github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRx
 github.com/gofrs/uuid v4.0.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gofrs/uuid v4.2.0+incompatible h1:yyYWMnhkhrKwwr8gAOcOCYxOOscHgDS9yZgBrnJfGa0=
 github.com/gofrs/uuid v4.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
-github.com/gogo/googleapis v0.0.0-20180223154316-0cd9801be74a/go.mod h1:gf4bu3Q80BeJ6H1S1vYPm8/ELATdvryBaNFGgqEef3s=
-github.com/gogo/googleapis v1.2.0/go.mod h1:Njal3psf3qN6dwBtQfUmBZh2ybovJ0tlu3o/AC7HYjU=
-github.com/gogo/googleapis v1.4.0/go.mod h1:5YRNX2z1oM5gXdAkurHa942MDgEJyk02w4OecKY87+c=
-github.com/gogo/googleapis v1.4.1 h1:1Yx4Myt7BxzvUr5ldGSbwYiZG6t9wGBZ+8/fX3Wvtq0=
-github.com/gogo/googleapis v1.4.1/go.mod h1:2lpHqI5OcWCtVElxXnPt+s8oJvMpySlOyM6xDCrzib4=
-github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
-github.com/gogo/protobuf v1.2.0/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
-github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
-github.com/gogo/protobuf v1.2.2-0.20190723190241-65acae22fc9d/go.mod h1:SlYgWuQ5SjCEi6WLHjHCa1yvBfUnHcTbrrZtXPKa29o=
-github.com/gogo/protobuf v1.3.0/go.mod h1:SlYgWuQ5SjCEi6WLHjHCa1yvBfUnHcTbrrZtXPKa29o=
-github.com/gogo/protobuf v1.3.1/go.mod h1:SlYgWuQ5SjCEi6WLHjHCa1yvBfUnHcTbrrZtXPKa29o=
-github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
-github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
-github.com/gogo/status v1.1.0 h1:+eIkrewn5q6b30y+g/BJINVVdi2xH7je5MPJ3ZPK3JA=
-github.com/gogo/status v1.1.0/go.mod h1:BFv9nrluPLmrS0EmGVvLaPNmRosr9KapBYd5/hpY1WM=
 github.com/golang-jwt/jwt/v4 v4.0.0/go.mod h1:/xlHOz8bRuivTWchD4jCa+NbatV+wEUSzwAxVc6locg=
 github.com/golang-jwt/jwt/v4 v4.1.0/go.mod h1:/xlHOz8bRuivTWchD4jCa+NbatV+wEUSzwAxVc6locg=
 github.com/golang-migrate/migrate/v4 v4.15.2 h1:vU+M05vs6jWHKDdmE1Ecwj0BznygFc4QsdRe2E/L7kc=

pkg/handler/handlerDstConn.go

@@ -6,10 +6,10 @@ import (
 	"net/http"
 	"strconv"
 
-	"github.com/gogo/status"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
 
 	"github.com/instill-ai/connector-backend/internal/resource"
 	"github.com/instill-ai/connector-backend/pkg/datamodel"

pkg/handler/handlerSrcConn.go

@@ -5,10 +5,10 @@ import (
 	"net/http"
 	"strconv"
 
-	"github.com/gogo/status"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
 
 	connectorPB "github.com/instill-ai/protogen-go/vdp/connector/v1alpha"
 )