From ba283695ddff22d990d37f6844ecbfdcfa9cd692 Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Sat, 1 May 2021 01:24:09 +0300
Subject: [PATCH] fix rsp code for wrappers

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 handler.go | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/handler.go b/handler.go
index 0c8d85b..d98f67b 100644
--- a/handler.go
+++ b/handler.go
@@ -205,20 +205,25 @@ func (h *httpServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		header:      md,
 	}
 
-	var scode int
 	// define the handler func
 	fn := func(fctx context.Context, req server.Request, rsp interface{}) (err error) {
 		returnValues = function.Call([]reflect.Value{hldr.rcvr, hldr.mtype.prepareContext(fctx), reflect.ValueOf(argv.Interface()), reflect.ValueOf(rsp)})
 
-		scode = GetRspCode(fctx)
 		// The return value for the method is an error.
 		if rerr := returnValues[0].Interface(); rerr != nil {
 			err = rerr.(error)
 		}
 
-		if md, ok := metadata.FromOutgoingContext(fctx); ok {
-			metadata.SetOutgoingContext(ctx, md)
+		md, ok := metadata.FromOutgoingContext(ctx)
+		if !ok {
+			md = metadata.New(0)
 		}
+		if nmd, ok := metadata.FromOutgoingContext(fctx); ok {
+			for k, v := range nmd {
+				md.Set(k, v)
+			}
+		}
+		metadata.SetOutgoingContext(ctx, md)
 
 		return err
 	}
@@ -237,6 +242,7 @@ func (h *httpServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		ct = DefaultContentType
 	}
 
+	scode := int(200)
 	if appErr := fn(ctx, hr, replyv.Interface()); appErr != nil {
 		switch verr := appErr.(type) {
 		case *errors.Error:
@@ -262,12 +268,12 @@ func (h *httpServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	w.Header().Set("content-Type", ct)
-	if scode != 0 {
-		w.WriteHeader(scode)
-	} else {
-		// handler.sopts.Logger.Warn(handler.sopts.Context, "response code not set in handler via SetRspCode(ctx, http.StatusXXX)")
-		w.WriteHeader(200)
+
+	if nscode := GetRspCode(ctx); nscode != 0 {
+		scode = nscode
 	}
+	w.WriteHeader(scode)
+
 	if _, cerr := w.Write(b); cerr != nil {
 		logger.DefaultLogger.Errorf(ctx, "write failed: %v", cerr)
 	}