<commit_before>package connection

import (
	"crypto/tls"
	"fmt"
	"github.com/fstab/h2c/http2client/frames"
	"github.com/fstab/h2c/http2client/internal/message"
	"github.com/fstab/h2c/http2client/internal/util"
	"golang.org/x/net/http2/hpack"
	"io"
	"net"
	"os"
)

const CLIENT_PREFACE = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"

// Some of these methods may no longer be needed after the last refactoring. Need to clean up.
type Connection interface {
	Host() string
	Port() int
	Error() error
	Shutdown()
	NewStream(request message.HttpRequest) Stream
	ServerFrameSize() uint32
	FetchPromisedStream(path string) Stream
	Write(stream Stream, frame frames.Frame)
	HandleIncomingFrame(frame frames.Frame)
	ReadNextFrame() (frames.Frame, error)
	HandleHttpRequest(request message.HttpRequest)
	HandleMonitoringRequest(request message.MonitoringRequest)
	HandlePingRequest(request message.PingRequest)
}

type connection struct {
	in                         chan frames.Frame
	shutdown                   chan bool
	info                       *info
	settings                   *settings
	streams                    map[uint32]*stream // StreamID -> *stream
	nextPingId                 uint64
	pendingPingRequests        map[uint64]message.PingRequest
	promisedStreamIDs          map[string]uint32 // Push Promise Path -> StreamID
	conn                       net.Conn
	isShutdown                 bool
	encodingContext            *frames.EncodingContext
	decodingContext            *frames.DecodingContext
	remainingSendWindowSize    int64
	remainingReceiveWindowSize int64
	incomingFrameFilters       []func(frames.Frame) frames.Frame
	outgoingFrameFilters       []func(frames.Frame) frames.Frame
	err                        error // TODO: not used
}

type info struct {
	host string
	port int
}

type settings struct {
	serverFrameSize                       uint32
	initialSendWindowSizeForNewStreams    uint32
	initialReceiveWindowSizeForNewStreams uint32
}

type writeFrameRequest struct {
	frame frames.Frame
	task  *util.AsyncTask
}

func Start(host string, port int, incomingFrameFilters []func(frames.Frame) frames.Frame, outgoingFrameFilters []func(frames.Frame) frames.Frame) (Connection, error) {
	hostAndPort := fmt.Sprintf("%v:%v", host, port)
	supportedProtocols := []string{"h2", "h2-16"} // The netty server still uses h2-16, treat it as if it was h2.
	conn, err := tls.Dial("tcp", hostAndPort, &tls.Config{
		InsecureSkipVerify: true,
		NextProtos:         supportedProtocols,
	})
	if err != nil {
		return nil, fmt.Errorf("Failed to connect to %v: %v", hostAndPort, err.Error())
	}
	if !util.SliceContainsString(supportedProtocols, conn.ConnectionState().NegotiatedProtocol) {
		return nil, fmt.Errorf("Server does not support HTTP/2 protocol.")
	}
	_, err = conn.Write([]byte(CLIENT_PREFACE))
	if err != nil {
		return nil, fmt.Errorf("Failed to write client preface to %v: %v", hostAndPort, err.Error())
	}
	c := newConnection(conn, host, port, incomingFrameFilters, outgoingFrameFilters)
	c.Write(nil, frames.NewSettingsFrame(0))
	return c, nil
}

func (conn *connection) HandleHttpRequest(request message.HttpRequest) {
	if conn.Error() != nil {
		request.CompleteWithError(conn.Error())
	}
	switch request.GetHeader(":method") {
	case "GET":
		conn.handleGetRequest(request)
	case "PUT":
		conn.handlePutRequest(request)
	case "POST":
		conn.handlePostRequest(request)
	default:
		request.CompleteWithError(fmt.Errorf("Request method '%v' not supported.", request.GetHeader(":method")))
	}
}

func (conn *connection) handleGetRequest(request message.HttpRequest) {
	stream := conn.FetchPromisedStream(request.GetHeader(":path"))
	if stream != nil {
		// Don't need to send request, because PUSH_PROMISE for this request already arrived.
		err := stream.AssociateWithRequest(request)
		if err != nil {
			request.CompleteWithError(err)
		}
	} else {
		conn.doRequest(request)
	}
}

func (conn *connection) handlePutRequest(request message.HttpRequest) {
	conn.doRequest(request)
}

func (conn *connection) handlePostRequest(request message.HttpRequest) {
	conn.doRequest(request)
}

func (conn *connection) doRequest(request message.HttpRequest) {
	stream := conn.NewStream(request)
	headersFrame := frames.NewHeadersFrame(stream.StreamId(), request.GetHeaders())
	headersFrame.EndStream = request.GetData() == nil
	conn.Write(stream, headersFrame)
	if request.GetData() != nil {
		conn.sendDataFrames(request.GetData(), stream)
	}
}

func (conn *connection) sendDataFrames(data []byte, stream Stream) {
	// chunkSize := uint32(len(data)) // use this to provoke GOAWAY frame with FRAME_SIZE_ERROR
	chunkSize := conn.ServerFrameSize() // TODO: Query chunk size with each iteration -> allow changes during loop
	nChunksSent := uint32(0)
	total := uint32(len(data))
	for nChunksSent*chunkSize < total {
		nextChunk := data[nChunksSent*chunkSize : min((nChunksSent+1)*chunkSize, total)]
		nChunksSent = nChunksSent + 1
		isLast := nChunksSent*chunkSize >= total
		dataFrame := frames.NewDataFrame(stream.StreamId(), nextChunk, isLast)
		conn.Write(stream, dataFrame)
	}
}

func min(a, b uint32) uint32 {
	if a < b {
		return a
	}
	return b
}

func (c *connection) HandleMonitoringRequest(request message.MonitoringRequest) {
	response := message.NewMonitoringResponse()
	for path := range c.promisedStreamIDs {
		response.AddPromisedPath(path)
	}
	request.CompleteSuccessfully(response)
}

func (c *connection) FetchPromisedStream(path string) Stream {
	streamId, exists := c.promisedStreamIDs[path]
	if exists {
		delete(c.promisedStreamIDs, path)
		return c.streams[streamId]
	} else {
		return nil
	}
}

func (c *connection) HandlePingRequest(request message.PingRequest) {
	pingFrame := frames.NewPingFrame(0, c.nextPingId, false)
	c.nextPingId = c.nextPingId + 1
	c.pendingPingRequests[pingFrame.Payload] = request
	c.writeImmediately(pingFrame)
}

func newConnection(conn net.Conn, host string, port int, incomingFrameFilters []func(frames.Frame) frames.Frame, outgoingFrameFilters []func(frames.Frame) frames.Frame) *connection {
	return &connection{
		in:       make(chan frames.Frame),
		shutdown: make(chan bool),
		info: &info{
			host: host,
			port: port,
		},
		settings: &settings{
			serverFrameSize:                       2 << 13,   // Minimum size that must be supported by all server implementations.
			initialSendWindowSizeForNewStreams:    2<<15 - 1, // Initial flow-control window size for new streams is 65,535 octets.
			initialReceiveWindowSizeForNewStreams: 2<<15 - 1,
		},
		streams:                    make(map[uint32]*stream),
		pendingPingRequests:        make(map[uint64]message.PingRequest),
		promisedStreamIDs:          make(map[string]uint32),
		isShutdown:                 false,
		conn:                       conn,
		encodingContext:            frames.NewEncodingContext(),
		decodingContext:            frames.NewDecodingContext(),
		remainingSendWindowSize:    2<<15 - 1,
		remainingReceiveWindowSize: 2<<15 - 1,
		incomingFrameFilters:       incomingFrameFilters,
		outgoingFrameFilters:       outgoingFrameFilters,
	}
}

func (c *connection) Shutdown() {
	c.isShutdown = true
	c.conn.Close()
}

func (c *connection) HandleIncomingFrame(frame frames.Frame) {
	switch frame := frame.(type) {
	case *frames.SettingsFrame:
		c.settings.handleSettingsFrame(frame)
	case *frames.HeadersFrame:
		stream := c.getOrCreateStream(frame.GetStreamId())
		stream.addReceivedHeaders(frame.Headers...)
		// TODO: continuations
		// TODO: error handling
		if frame.EndStream {
			stream.endStream()
		}
	case *frames.DataFrame:
		stream, exists := c.streams[frame.GetStreamId()]
		if !exists {
			// TODO: error handling
			fmt.Fprintf(os.Stderr, "Received data for unknown stream %v. Ignoring this frame.", frame.GetStreamId())
			return
		}
		c.flowControlForIncomingDataFrame(frame, stream)
		stream.appendReceivedData(frame.Data)
		if frame.EndStream {
			stream.endStream()
		}
	case *frames.PushPromiseFrame:
		method := findHeader(":method", frame.Headers)
		path := findHeader(":path", frame.Headers)
		if method != "GET" {
			fmt.Fprintf(os.Stderr, "ERROR: PUSH_PROMISE with method %v not supported.", method)
			return
		}
		_, exists := c.streams[frame.PromisedStreamId]
		if exists {
			fmt.Fprintf(os.Stderr, "ERROR: Received PUSH_PROMISE for existing stream %v", frame.PromisedStreamId)
			return
		}
		if !frame.EndHeaders {
			fmt.Fprintf(os.Stderr, "ERROR: Push promise with multiple header frames not supported.")
			return
		}
		c.getOrCreateStream(frame.PromisedStreamId)
		c.promisedStreamIDs[path] = frame.PromisedStreamId
	case *frames.PingFrame:
		if frame.Ack {
			pendingPingRequest, exists := c.pendingPingRequests[frame.Payload]
			if exists {
				delete(c.pendingPingRequests, frame.Payload)
				pendingPingRequest.CompleteSuccessfully(message.NewPingResponse())
			} else {
				pingFrame := frames.NewPingFrame(0, frame.Payload, true)
				c.writeImmediately(pingFrame)
			}
		}
	case *frames.RstStreamFrame:
		stream, exists := c.streams[frame.GetStreamId()]
		if !exists {
			// TODO: error handling
			fmt.Fprintf(os.Stderr, "Received data for unknown stream %v. Ignoring this frame.\n", frame.GetStreamId())
			return
		}
		stream.setError(fmt.Errorf("ERROR: Server sent RST_STREAM with error code %v.", frame.ErrorCode.String()))
		stream.endStream()
	case *frames.WindowUpdateFrame:
		c.handleWindowUpdateFrame(frame)
	case *frames.GoAwayFrame:
		// TODO: error handling
		fmt.Fprintf(os.Stderr, "Connection closed: Server sent GOAWAY with error code %v\n", frame.ErrorCode.String())
		c.shutdown <- true
	default:
		// TODO: error handling
		fmt.Fprintf(os.Stderr, "Received unknown frame type %v\n", frame.Type())
	}
}<commit_msg>fix bug when receiving ping frames<commit_after>package connection

import (
	"crypto/tls"
	"fmt"
	"github.com/fstab/h2c/http2client/frames"
	"github.com/fstab/h2c/http2client/internal/message"
	"github.com/fstab/h2c/http2client/internal/util"
	"golang.org/x/net/http2/hpack"
	"io"
	"net"
	"os"
)

const CLIENT_PREFACE = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"

// Some of these methods may no longer be needed after the last refactoring. Need to clean up.
type Connection interface {
	Host() string
	Port() int
	Error() error
	Shutdown()
	NewStream(request message.HttpRequest) Stream
	ServerFrameSize() uint32
	FetchPromisedStream(path string) Stream
	Write(stream Stream, frame frames.Frame)
	HandleIncomingFrame(frame frames.Frame)
	ReadNextFrame() (frames.Frame, error)
	HandleHttpRequest(request message.HttpRequest)
	HandleMonitoringRequest(request message.MonitoringRequest)
	HandlePingRequest(request message.PingRequest)
}

type connection struct {
	in                         chan frames.Frame
	shutdown                   chan bool
	info                       *info
	settings                   *settings
	streams                    map[uint32]*stream // StreamID -> *stream
	nextPingId                 uint64
	pendingPingRequests        map[uint64]message.PingRequest
	promisedStreamIDs          map[string]uint32 // Push Promise Path -> StreamID
	conn                       net.Conn
	isShutdown                 bool
	encodingContext            *frames.EncodingContext
	decodingContext            *frames.DecodingContext
	remainingSendWindowSize    int64
	remainingReceiveWindowSize int64
	incomingFrameFilters       []func(frames.Frame) frames.Frame
	outgoingFrameFilters       []func(frames.Frame) frames.Frame
	err                        error // TODO: not used
}

type info struct {
	host string
	port int
}

type settings struct {
	serverFrameSize                       uint32
	initialSendWindowSizeForNewStreams    uint32
	initialReceiveWindowSizeForNewStreams uint32
}

type writeFrameRequest struct {
	frame frames.Frame
	task  *util.AsyncTask
}

func Start(host string, port int, incomingFrameFilters []func(frames.Frame) frames.Frame, outgoingFrameFilters []func(frames.Frame) frames.Frame) (Connection, error) {
	hostAndPort := fmt.Sprintf("%v:%v", host, port)
	supportedProtocols := []string{"h2", "h2-16"} // The netty server still uses h2-16, treat it as if it was h2.
	conn, err := tls.Dial("tcp", hostAndPort, &tls.Config{
		InsecureSkipVerify: true,
		NextProtos:         supportedProtocols,
	})
	if err != nil {
		return nil, fmt.Errorf("Failed to connect to %v: %v", hostAndPort, err.Error())
	}
	if !util.SliceContainsString(supportedProtocols, conn.ConnectionState().NegotiatedProtocol) {
		return nil, fmt.Errorf("Server does not support HTTP/2 protocol.")
	}
	_, err = conn.Write([]byte(CLIENT_PREFACE))
	if err != nil {
		return nil, fmt.Errorf("Failed to write client preface to %v: %v", hostAndPort, err.Error())
	}
	c := newConnection(conn, host, port, incomingFrameFilters, outgoingFrameFilters)
	c.Write(nil, frames.NewSettingsFrame(0))
	return c, nil
}

func (conn *connection) HandleHttpRequest(request message.HttpRequest) {
	if conn.Error() != nil {
		request.CompleteWithError(conn.Error())
	}
	switch request.GetHeader(":method") {
	case "GET":
		conn.handleGetRequest(request)
	case "PUT":
		conn.handlePutRequest(request)
	case "POST":
		conn.handlePostRequest(request)
	default:
		request.CompleteWithError(fmt.Errorf("Request method '%v' not supported.", request.GetHeader(":method")))
	}
}

func (conn *connection) handleGetRequest(request message.HttpRequest) {
	stream := conn.FetchPromisedStream(request.GetHeader(":path"))
	if stream != nil {
		// Don't need to send request, because PUSH_PROMISE for this request already arrived.
		err := stream.AssociateWithRequest(request)
		if err != nil {
			request.CompleteWithError(err)
		}
	} else {
		conn.doRequest(request)
	}
}

func (conn *connection) handlePutRequest(request message.HttpRequest) {
	conn.doRequest(request)
}

func (conn *connection) handlePostRequest(request message.HttpRequest) {
	conn.doRequest(request)
}

func (conn *connection) doRequest(request message.HttpRequest) {
	stream := conn.NewStream(request)
	headersFrame := frames.NewHeadersFrame(stream.StreamId(), request.GetHeaders())
	headersFrame.EndStream = request.GetData() == nil
	conn.Write(stream, headersFrame)
	if request.GetData() != nil {
		conn.sendDataFrames(request.GetData(), stream)
	}
}

func (conn *connection) sendDataFrames(data []byte, stream Stream) {
	// chunkSize := uint32(len(data)) // use this to provoke GOAWAY frame with FRAME_SIZE_ERROR
	chunkSize := conn.ServerFrameSize() // TODO: Query chunk size with each iteration -> allow changes during loop
	nChunksSent := uint32(0)
	total := uint32(len(data))
	for nChunksSent*chunkSize < total {
		nextChunk := data[nChunksSent*chunkSize : min((nChunksSent+1)*chunkSize, total)]
		nChunksSent = nChunksSent + 1
		isLast := nChunksSent*chunkSize >= total
		dataFrame := frames.NewDataFrame(stream.StreamId(), nextChunk, isLast)
		conn.Write(stream, dataFrame)
	}
}

func min(a, b uint32) uint32 {
	if a < b {
		return a
	}
	return b
}

func (c *connection) HandleMonitoringRequest(request message.MonitoringRequest) {
	response := message.NewMonitoringResponse()
	for path := range c.promisedStreamIDs {
		response.AddPromisedPath(path)
	}
	request.CompleteSuccessfully(response)
}

func (c *connection) FetchPromisedStream(path string) Stream {
	streamId, exists := c.promisedStreamIDs[path]
	if exists {
		delete(c.promisedStreamIDs, path)
		return c.streams[streamId]
	} else {
		return nil
	}
}

func (c *connection) HandlePingRequest(request message.PingRequest) {
	pingFrame := frames.NewPingFrame(0, c.nextPingId, false)
	c.nextPingId = c.nextPingId + 1
	c.pendingPingRequests[pingFrame.Payload] = request
	c.writeImmediately(pingFrame)
}

func newConnection(conn net.Conn, host string, port int, incomingFrameFilters []func(frames.Frame) frames.Frame, outgoingFrameFilters []func(frames.Frame) frames.Frame) *connection {
	return &connection{
		in:       make(chan frames.Frame),
		shutdown: make(chan bool),
		info: &info{
			host: host,
			port: port,
		},
		settings: &settings{
			serverFrameSize:                       2 << 13,   // Minimum size that must be supported by all server implementations.
			initialSendWindowSizeForNewStreams:    2<<15 - 1, // Initial flow-control window size for new streams is 65,535 octets.
			initialReceiveWindowSizeForNewStreams: 2<<15 - 1,
		},
		streams:                    make(map[uint32]*stream),
		pendingPingRequests:        make(map[uint64]message.PingRequest),
		promisedStreamIDs:          make(map[string]uint32),
		isShutdown:                 false,
		conn:                       conn,
		encodingContext:            frames.NewEncodingContext(),
		decodingContext:            frames.NewDecodingContext(),
		remainingSendWindowSize:    2<<15 - 1,
		remainingReceiveWindowSize: 2<<15 - 1,
		incomingFrameFilters:       incomingFrameFilters,
		outgoingFrameFilters:       outgoingFrameFilters,
	}
}

func (c *connection) Shutdown() {
	c.isShutdown = true
	c.conn.Close()
}

func (c *connection) HandleIncomingFrame(frame frames.Frame) {
	switch frame := frame.(type) {
	case *frames.SettingsFrame:
		c.settings.handleSettingsFrame(frame)
	case *frames.HeadersFrame:
		stream := c.getOrCreateStream(frame.GetStreamId())
		stream.addReceivedHeaders(frame.Headers...)
		// TODO: continuations
		// TODO: error handling
		if frame.EndStream {
			stream.endStream()
		}
	case *frames.DataFrame:
		stream, exists := c.streams[frame.GetStreamId()]
		if !exists {
			// TODO: error handling
			fmt.Fprintf(os.Stderr, "Received data for unknown stream %v. Ignoring this frame.", frame.GetStreamId())
			return
		}
		c.flowControlForIncomingDataFrame(frame, stream)
		stream.appendReceivedData(frame.Data)
		if frame.EndStream {
			stream.endStream()
		}
	case *frames.PushPromiseFrame:
		method := findHeader(":method", frame.Headers)
		path := findHeader(":path", frame.Headers)
		if method != "GET" {
			fmt.Fprintf(os.Stderr, "ERROR: PUSH_PROMISE with method %v not supported.", method)
			return
		}
		_, exists := c.streams[frame.PromisedStreamId]
		if exists {
			fmt.Fprintf(os.Stderr, "ERROR: Received PUSH_PROMISE for existing stream %v", frame.PromisedStreamId)
			return
		}
		if !frame.EndHeaders {
			fmt.Fprintf(os.Stderr, "ERROR: Push promise with multiple header frames not supported.")
			return
		}
		c.getOrCreateStream(frame.PromisedStreamId)
		c.promisedStreamIDs[path] = frame.PromisedStreamId
	case *frames.PingFrame:
		if frame.Ack {
			pendingPingRequest, exists := c.pendingPingRequests[frame.Payload]
			if exists {
				delete(c.pendingPingRequests, frame.Payload)
				pendingPingRequest.CompleteSuccessfully(message.NewPingResponse())
			}
		} else {
			pingFrame := frames.NewPingFrame(0, frame.Payload, true)
			c.writeImmediately(pingFrame)
		}
	case *frames.RstStreamFrame:
		stream, exists := c.streams[frame.GetStreamId()]
		if !exists {
			// TODO: error handling
			fmt.Fprintf(os.Stderr, "Received data for unknown stream %v. Ignoring this frame.\n", frame.GetStreamId())
			return
		}
		stream.setError(fmt.Errorf("ERROR: Server sent RST_STREAM with error code %v.", frame.ErrorCode.String()))
		stream.endStream()
	case *frames.WindowUpdateFrame:
		c.handleWindowUpdateFrame(frame)
	case *frames.GoAwayFrame:
		// TODO: error handling
		fmt.Fprintf(os.Stderr, "Connection closed: Server sent GOAWAY with error code %v\n", frame.ErrorCode.String())
		c.shutdown <- true
	default:
		// TODO: error handling
		fmt.Fprintf(os.Stderr, "Received unknown frame type %v\n", frame.Type())
	}
}
{'fraction_non_alphanumeric': 0.11290821056969197, 'fraction_numerical': 0.008756567425569177, 'mean_word_length': 3.9794819184406256, 'pattern_counts': {'":': 16, '<': 29, '<?xml version=': 0, '>': 11, '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': '2358220', 'n_tokens_mistral': 6494, 'n_tokens_neox': 6104, 'n_words': 1421}