<commit_before>package writer

import (
	"bufio"
	"io"
	"net"
	"net/http"
)

// NewResponseWriter wraps an http.ResponseWriter, returning a proxy that allows you to
// hook into various parts of the response process.
func NewResponseWriter(w http.ResponseWriter, protoMajor int) ResponseWriter {
	bw := BasicWriter{inner: w}
	if protoMajor == 2 {
		return &HttpTwoWriter{bw}
	}
	return &HttpOneWriter{bw}
}

func SetProxyStatusCode(w ResponseWriter, code int) {
	w.(*BasicWriter).code = code
}

// ResponseWriter is a proxy around an http.ResponseWriter that allows you to hook
// into various parts of the response process.
type ResponseWriter interface {
	http.ResponseWriter
	// Status returns the HTTP status of the request, or 0 if one has not
	// yet been sent.
	Status() int
	// BytesWritten returns the total number of bytes sent to the client.
	BytesWritten() int
	// Tee causes the response body to be written to the given io.Writer in
	// addition to proxying the writes through. Only one io.Writer can be
	// tee'd to at once: setting a second one will overwrite the first.
	// Writes will be sent to the proxy before being written to this
	// io.Writer. It is illegal for the tee'd writer to be modified
	// concurrently with writes.
	Tee(io.Writer)
	// Unwrap returns the original proxied target.
	Unwrap() http.ResponseWriter
	Flush()
}

// basicWriter wraps a http.ResponseWriter that implements the minimal
// http.ResponseWriter interface.
type BasicWriter struct {
	inner         http.ResponseWriter
	wroteHeader   bool
	headerFlushed bool
	code          int
	bytes         int
	tee           io.Writer
}

func (b *BasicWriter) Header() http.Header {
	return b.inner.Header()
}

func (b *BasicWriter) WriteHeader(code int) {
	b.code = code
	if code != 0 {
		b.wroteHeader = true
	} else {
		b.wroteHeader = false
	}
}

func (b *BasicWriter) WriteHeaderImmediate(code int) {
	b.WriteHeader(code)
	if b.wroteHeader {
		b.inner.WriteHeader(code)
		b.headerFlushed = true
	}
}

func (b *BasicWriter) FlushHeadersIfRequired() {
	if !b.headerFlushed {
		b.WriteHeaderImmediate(b.code)
	}
}

func (b *BasicWriter) Write(buf []byte) (int, error) {
	b.FlushHeadersIfRequired()
	n, err := b.inner.Write(buf)
	if b.tee != nil {
		_, err2 := b.tee.Write(buf)
		// Prefer errors generated by the proxied writer.
		if err == nil {
			err = err2
		}
	}
	b.bytes += n
	return n, err
}

func (b *BasicWriter) Status() int {
	// if b.code == 0 {
	// 	return 200
	// }
	return b.code
}

func (b *BasicWriter) BytesWritten() int {
	return b.bytes
}

func (b *BasicWriter) Tee(w io.Writer) {
	b.tee = w
}

func (b *BasicWriter) Unwrap() http.ResponseWriter {
	return b.inner
}

func (b *BasicWriter) CloseNotify() <-chan bool {
	cn := b.inner.(http.CloseNotifier)
	return cn.CloseNotify()
}

func (b *BasicWriter) Flush() {
	b.FlushHeadersIfRequired()
	fl := b.inner.(http.Flusher)
	fl.Flush()
}

// HttpOneWriter is a HTTP writer that additionally satisfies http.Hijacker,
// and io.ReaderFrom.
type HttpOneWriter struct {
	BasicWriter
}

func (f *HttpOneWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hj := f.BasicWriter.inner.(http.Hijacker)
	return hj.Hijack()
}

func (f *HttpOneWriter) ReadFrom(r io.Reader) (int64, error) {
	f.BasicWriter.FlushHeadersIfRequired()
	if f.BasicWriter.tee != nil {
		return io.Copy(&f.BasicWriter, r)
	}
	rf := f.BasicWriter.inner.(io.ReaderFrom)
	n, err := rf.ReadFrom(r)
	f.BasicWriter.bytes += int(n)
	return n, err
}

// HttpTwoWriter is a HTTP2 writer that additionally satisfies
// Push
type HttpTwoWriter struct {
	BasicWriter
}

func (f *HttpTwoWriter) Push(target string, opts *http.PushOptions) error {
	return f.BasicWriter.inner.(http.Pusher).Push(target, opts)
}<commit_msg>add: IsHijacked to ResponseWriter<commit_after>package writer

import (
	"bufio"
	"io"
	"net"
	"net/http"
)

// NewResponseWriter wraps an http.ResponseWriter, returning a proxy that allows you to
// hook into various parts of the response process.
func NewResponseWriter(w http.ResponseWriter, protoMajor int) ResponseWriter {
	bw := BasicWriter{inner: w}
	if protoMajor == 2 {
		return &HttpTwoWriter{bw}
	}
	return &HttpOneWriter{bw, false}
}

// ResponseWriter is a proxy around an http.ResponseWriter that allows you to hook
// into various parts of the response process.
type ResponseWriter interface {
	http.ResponseWriter
	// Status returns the HTTP status of the request, or 0 if one has not
	// yet been sent.
	Status() int
	// BytesWritten returns the total number of bytes sent to the client.
	BytesWritten() int
	// Tee causes the response body to be written to the given io.Writer in
	// addition to proxying the writes through. Only one io.Writer can be
	// tee'd to at once: setting a second one will overwrite the first.
	// Writes will be sent to the proxy before being written to this
	// io.Writer. It is illegal for the tee'd writer to be modified
	// concurrently with writes.
	Tee(io.Writer)
	// Unwrap returns the original proxied target.
	Unwrap() http.ResponseWriter
	Flush()
	IsHijacked() bool
}

// basicWriter wraps a http.ResponseWriter that implements the minimal
// http.ResponseWriter interface.
type BasicWriter struct {
	inner           http.ResponseWriter
	isHeaderSet     bool
	isHeaderFlushed bool
	code            int
	bytes           int
	tee             io.Writer
}

func (b *BasicWriter) Header() http.Header {
	return b.inner.Header()
}

func (b *BasicWriter) WriteHeader(code int) {
	b.code = code
	if code != 0 {
		b.isHeaderSet = true
	} else {
		b.isHeaderSet = false
	}
}

func (b *BasicWriter) WriteHeaderImmediate(code int) {
	b.WriteHeader(code)
	if b.isHeaderSet {
		b.inner.WriteHeader(code)
		b.isHeaderFlushed = true
	}
}

func (b *BasicWriter) FlushHeadersIfRequired() {
	if !b.isHeaderFlushed {
		b.WriteHeaderImmediate(b.code)
	}
}

func (b *BasicWriter) Write(buf []byte) (int, error) {
	b.FlushHeadersIfRequired()
	n, err := b.inner.Write(buf)
	if b.tee != nil {
		_, err2 := b.tee.Write(buf)
		// Prefer errors generated by the proxied writer.
		if err == nil {
			err = err2
		}
	}
	b.bytes += n
	return n, err
}

func (b *BasicWriter) Status() int {
	// if b.code == 0 {
	// 	return 200
	// }
	return b.code
}

func (b *BasicWriter) BytesWritten() int {
	return b.bytes
}

func (b *BasicWriter) Tee(w io.Writer) {
	b.tee = w
}

func (b *BasicWriter) Unwrap() http.ResponseWriter {
	return b.inner
}

func (b *BasicWriter) CloseNotify() <-chan bool {
	cn := b.inner.(http.CloseNotifier)
	return cn.CloseNotify()
}

func (b *BasicWriter) Flush() {
	b.FlushHeadersIfRequired()
	fl := b.inner.(http.Flusher)
	fl.Flush()
}

func (b *BasicWriter) IsHijacked() bool {
	return false
}

// HttpOneWriter is a HTTP writer that additionally satisfies http.Hijacker,
// and io.ReaderFrom.
type HttpOneWriter struct {
	BasicWriter
	isHijacked bool
}

func (f *HttpOneWriter) IsHijacked() bool {
	return f.isHijacked
}

func (f *HttpOneWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hj := f.BasicWriter.inner.(http.Hijacker)
	conn, rw, err := hj.Hijack()
	if err == nil {
		f.isHijacked = true
		return conn, rw, err
	}
	return conn, rw, err
}

func (f *HttpOneWriter) ReadFrom(r io.Reader) (int64, error) {
	f.BasicWriter.FlushHeadersIfRequired()
	if f.BasicWriter.tee != nil {
		return io.Copy(&f.BasicWriter, r)
	}
	rf := f.BasicWriter.inner.(io.ReaderFrom)
	n, err := rf.ReadFrom(r)
	f.BasicWriter.bytes += int(n)
	return n, err
}

// HttpTwoWriter is a HTTP2 writer that additionally satisfies
// Push
type HttpTwoWriter struct {
	BasicWriter
}

func (f *HttpTwoWriter) Push(target string, opts *http.PushOptions) error {
	return f.BasicWriter.inner.(http.Pusher).Push(target, opts)
}
{'fraction_non_alphanumeric': 0.11328125, 'fraction_numerical': 0.003125, 'mean_word_length': 3.6104441776710683, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6097129', 'n_tokens_mistral': 2693, 'n_tokens_neox': 2610, 'n_words': 828}