- Initial broadcast commit

Signed-off-by: Valery Piashchynski <[email protected]>

1 files changed, 105 insertions, 0 deletions

diff --git a/pkg/worker_handler/response.go b/pkg/worker_handler/response.go
new file mode 100644
index 00000000..1763d304
--- /dev/null
+++ b/pkg/worker_handler/response.go
@@ -0,0 +1,105 @@
+package handler
+
+import (
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+
+	"github.com/spiral/roadrunner/v2/pkg/payload"
+)
+
+// Response handles PSR7 response logic.
+type Response struct {
+	// Status contains response status.
+	Status int `json:"status"`
+
+	// Header contains list of response headers.
+	Headers map[string][]string `json:"headers"`
+
+	// associated Body payload.
+	Body interface{}
+	sync.Mutex
+}
+
+// NewResponse creates new response based on given pool payload.
+func NewResponse(p payload.Payload) (*Response, error) {
+	r := &Response{Body: p.Body}
+	if err := json.Unmarshal(p.Context, r); err != nil {
+		return nil, err
+	}
+
+	return r, nil
+}
+
+// Write writes response headers, status and body into ResponseWriter.
+func (r *Response) Write(w http.ResponseWriter) error {
+	// INFO map is the reference type in golang
+	p := handlePushHeaders(r.Headers)
+	if pusher, ok := w.(http.Pusher); ok {
+		for _, v := range p {
+			err := pusher.Push(v, nil)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	handleTrailers(r.Headers)
+	for n, h := range r.Headers {
+		for _, v := range h {
+			w.Header().Add(n, v)
+		}
+	}
+
+	w.WriteHeader(r.Status)
+
+	if data, ok := r.Body.([]byte); ok {
+		_, err := w.Write(data)
+		if err != nil {
+			return handleWriteError(err)
+		}
+	}
+
+	if rc, ok := r.Body.(io.Reader); ok {
+		if _, err := io.Copy(w, rc); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func handlePushHeaders(h map[string][]string) []string {
+	var p []string
+	pushHeader, ok := h[http2pushHeaderKey]
+	if !ok {
+		return p
+	}
+
+	p = append(p, pushHeader...)
+
+	delete(h, http2pushHeaderKey)
+
+	return p
+}
+
+func handleTrailers(h map[string][]string) {
+	trailers, ok := h[TrailerHeaderKey]
+	if !ok {
+		return
+	}
+
+	for _, tr := range trailers {
+		for _, n := range strings.Split(tr, ",") {
+			n = strings.Trim(n, "\t ")
+			if v, ok := h[n]; ok {
+				h["Trailer:"+n] = v
+
+				delete(h, n)
+			}
+		}
+	}
+
+	delete(h, TrailerHeaderKey)
+}