github.com/Prakhar-Agarwal-byte/moby@v0.0.0-20231027092010-a14e3e8ab87e/pkg/authorization/response.go

package authorization // import "github.com/Prakhar-Agarwal-byte/moby/pkg/authorization"

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"

	"github.com/containerd/log"
)

// ResponseModifier allows authorization plugins to read and modify the content of the http.response
type ResponseModifier interface {
	http.ResponseWriter
	http.Flusher

	// RawBody returns the current http content
	RawBody() []byte

	// RawHeaders returns the current content of the http headers
	RawHeaders() ([]byte, error)

	// StatusCode returns the current status code
	StatusCode() int

	// OverrideBody replaces the body of the HTTP reply
	OverrideBody(b []byte)

	// OverrideHeader replaces the headers of the HTTP reply
	OverrideHeader(b []byte) error

	// OverrideStatusCode replaces the status code of the HTTP reply
	OverrideStatusCode(statusCode int)

	// FlushAll flushes all data to the HTTP response
	FlushAll() error

	// Hijacked indicates the response has been hijacked by the Docker daemon
	Hijacked() bool
}

// NewResponseModifier creates a wrapper to an http.ResponseWriter to allow inspecting and modifying the content
func NewResponseModifier(rw http.ResponseWriter) ResponseModifier {
	return &responseModifier{rw: rw, header: make(http.Header)}
}

const maxBufferSize = 64 * 1024

// responseModifier is used as an adapter to http.ResponseWriter in order to manipulate and explore
// the http request/response from docker daemon
type responseModifier struct {
	// The original response writer
	rw http.ResponseWriter
	// body holds the response body
	body []byte
	// header holds the response header
	header http.Header
	// statusCode holds the response status code
	statusCode int
	// hijacked indicates the request has been hijacked
	hijacked bool
}

func (rm *responseModifier) Hijacked() bool {
	return rm.hijacked
}

// WriteHeader stores the http status code
func (rm *responseModifier) WriteHeader(s int) {
	// Use original request if hijacked
	if rm.hijacked {
		rm.rw.WriteHeader(s)
		return
	}

	rm.statusCode = s
}

// Header returns the internal http header
func (rm *responseModifier) Header() http.Header {
	// Use original header if hijacked
	if rm.hijacked {
		return rm.rw.Header()
	}

	return rm.header
}

// StatusCode returns the http status code
func (rm *responseModifier) StatusCode() int {
	return rm.statusCode
}

// OverrideBody replaces the body of the HTTP response
func (rm *responseModifier) OverrideBody(b []byte) {
	rm.body = b
}

// OverrideStatusCode replaces the status code of the HTTP response
func (rm *responseModifier) OverrideStatusCode(statusCode int) {
	rm.statusCode = statusCode
}

// OverrideHeader replaces the headers of the HTTP response
func (rm *responseModifier) OverrideHeader(b []byte) error {
	header := http.Header{}
	if err := json.Unmarshal(b, &header); err != nil {
		return err
	}
	rm.header = header
	return nil
}

// Write stores the byte array inside content
func (rm *responseModifier) Write(b []byte) (int, error) {
	if rm.hijacked {
		return rm.rw.Write(b)
	}

	if len(rm.body)+len(b) > maxBufferSize {
		rm.Flush()
	}
	rm.body = append(rm.body, b...)
	return len(b), nil
}

// Body returns the response body
func (rm *responseModifier) RawBody() []byte {
	return rm.body
}

func (rm *responseModifier) RawHeaders() ([]byte, error) {
	var b bytes.Buffer
	if err := rm.header.Write(&b); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

// Hijack returns the internal connection of the wrapped http.ResponseWriter
func (rm *responseModifier) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	rm.hijacked = true
	rm.FlushAll()

	hijacker, ok := rm.rw.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("Internal response writer doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}

// Flush uses the internal flush API of the wrapped http.ResponseWriter
func (rm *responseModifier) Flush() {
	flusher, ok := rm.rw.(http.Flusher)
	if !ok {
		log.G(context.TODO()).Error("Internal response writer doesn't support the Flusher interface")
		return
	}

	rm.FlushAll()
	flusher.Flush()
}

// FlushAll flushes all data to the HTTP response
func (rm *responseModifier) FlushAll() error {
	// Copy the header
	for k, vv := range rm.header {
		for _, v := range vv {
			rm.rw.Header().Add(k, v)
		}
	}

	// Copy the status code
	// Also WriteHeader needs to be done after all the headers
	// have been copied (above).
	if rm.statusCode > 0 {
		rm.rw.WriteHeader(rm.statusCode)
	}

	var err error
	if len(rm.body) > 0 {
		// Write body
		var n int
		n, err = rm.rw.Write(rm.body)
		// TODO(@cpuguy83): there is now a relatively small buffer limit, instead of discarding our buffer here and
		// allocating again later this should just keep using the same buffer and track the buffer position (like a bytes.Buffer with a fixed size)
		rm.body = rm.body[n:]
	}

	// Clean previous data
	rm.statusCode = 0
	rm.header = http.Header{}
	return err
}