// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package http3

import (
	"bytes"
	"errors"
	"io"
	"net/http"
	"net/http/httptrace"
	"net/textproto"
	"reflect"
	"slices"
	"strings"
	"testing"
	"testing/synctest"

	"golang.org/x/net/quic"
)

func TestRoundTripSimple(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		req, _ := http.NewRequest("GET", "https://example.tld/", nil)
		req.Header["User-Agent"] = nil
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(http.Header{
			":authority": []string{"example.tld"},
			":method":    []string{"GET"},
			":path":      []string{"/"},
			":scheme":    []string{"https"},
		})
		st.writeHeaders(http.Header{
			":status":       []string{"200"},
			"x-some-header": []string{"value"},
		})
		rt.wantStatus(200)
		rt.wantHeaders(http.Header{
			"X-Some-Header": []string{"value"},
		})
	})
}

func TestRoundTripWithBadHeaders(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		req, _ := http.NewRequest("GET", "https://example.tld/", nil)
		req.Header["Invalid\nHeader"] = []string{"x"}
		rt := tc.roundTrip(req)
		rt.wantError("RoundTrip fails when request contains invalid headers")
	})
}

func TestRoundTripWithUnknownFrame(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		req, _ := http.NewRequest("GET", "https://example.tld/", nil)
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)

		// Write an unknown frame type before the response HEADERS.
		data := "frame content"
		st.writeVarint(0x1f + 0x21)      // reserved frame type
		st.writeVarint(int64(len(data))) // size
		st.Write([]byte(data))

		st.writeHeaders(http.Header{
			":status": []string{"200"},
		})
		rt.wantStatus(200)
	})
}

func TestRoundTripWithInvalidPushPromise(t *testing.T) {
	// "A client MUST treat receipt of a PUSH_PROMISE frame that contains
	// a larger push ID than the client has advertised as a connection error of H3_ID_ERROR."
	// https://www.rfc-editor.org/rfc/rfc9114.html#section-7.2.5-5
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		req, _ := http.NewRequest("GET", "https://example.tld/", nil)
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)

		// Write a PUSH_PROMISE frame.
		// Since the client hasn't indicated willingness to accept pushes,
		// this is a connection error.
		st.writePushPromise(0, http.Header{
			":path": []string{"/foo"},
		})
		rt.wantError("RoundTrip fails after receiving invalid PUSH_PROMISE")
		tc.wantClosed(
			"push ID exceeds client's MAX_PUSH_ID",
			errH3IDError,
		)
	})
}

func TestRoundTripResponseContentLength(t *testing.T) {
	for _, test := range []struct {
		name              string
		respHeader        http.Header
		wantContentLength int64
		wantError         bool
	}{{
		name: "valid",
		respHeader: http.Header{
			":status":        []string{"200"},
			"content-length": []string{"100"},
		},
		wantContentLength: 100,
	}, {
		name: "absent",
		respHeader: http.Header{
			":status": []string{"200"},
		},
		wantContentLength: -1,
	}, {
		name: "unparsable",
		respHeader: http.Header{
			":status":        []string{"200"},
			"content-length": []string{"1 1"},
		},
		wantError: true,
	}, {
		name: "duplicated",
		respHeader: http.Header{
			":status":        []string{"200"},
			"content-length": []string{"500", "500", "500"},
		},
		wantContentLength: 500,
	}, {
		name: "inconsistent",
		respHeader: http.Header{
			":status":        []string{"200"},
			"content-length": []string{"1", "2"},
		},
		wantError: true,
	}, {
		// 204 responses aren't allowed to contain a Content-Length header.
		// We just ignore it.
		name: "204",
		respHeader: http.Header{
			":status":        []string{"204"},
			"content-length": []string{"100"},
		},
		wantContentLength: -1,
	}} {
		synctestSubtest(t, test.name, func(t *testing.T) {
			tc := newTestClientConn(t)
			tc.greet()

			req, _ := http.NewRequest("GET", "https://example.tld/", nil)
			rt := tc.roundTrip(req)
			st := tc.wantStream(streamTypeRequest)
			st.wantHeaders(nil)
			st.writeHeaders(test.respHeader)
			if test.wantError {
				rt.wantError("invalid content-length in response")
				return
			}
			if got, want := rt.response().ContentLength, test.wantContentLength; got != want {
				t.Errorf("Response.ContentLength = %v, want %v", got, want)
			}
		})
	}
}

func TestRoundTripMalformedResponses(t *testing.T) {
	for _, test := range []struct {
		name       string
		respHeader http.Header
	}{{
		name: "duplicate :status",
		respHeader: http.Header{
			":status": {"200", "204"},
		},
	}, {
		name: "unparsable :status",
		respHeader: http.Header{
			":status": {"frogpants"},
		},
	}, {
		name: "undefined pseudo-header",
		respHeader: http.Header{
			":status":  {"200"},
			":unknown": {"x"},
		},
	}, {
		name:       "no :status",
		respHeader: http.Header{},
	}, {
		name: "header name with control character",
		respHeader: http.Header{
			":status":             {"200"},
			"name\nevilinjection": {"Value"},
		},
	}, {
		name: "header name with uppercase character",
		respHeader: http.Header{
			":status": {"200"},
			"nAme":    {"Value"},
		},
	}, {
		name:       "pseudo-header name with control character",
		respHeader: http.Header{":status\nevilinjection": {"200"}},
	}, {
		name:       "pseudo-header name with uppercase character",
		respHeader: http.Header{":stAtus": {"200"}},
	}, {
		name: "header value with control character",
		respHeader: http.Header{
			":status": {"200"},
			"name":    {"Value\nEvilInjection"},
		},
	}, {
		name:       "pseudo-header value with control character",
		respHeader: http.Header{":status": {"200\nEvilInjection"}},
	}} {
		synctestSubtest(t, test.name, func(t *testing.T) {
			tc := newTestClientConn(t)
			tc.greet()

			req, _ := http.NewRequest("GET", "https://example.tld/", nil)
			rt := tc.roundTrip(req)
			st := tc.wantStream(streamTypeRequest)
			st.wantHeaders(nil)
			st.writeHeadersRaw(test.respHeader)
			rt.wantError("malformed response")
		})
	}
}

func TestRoundTripCrumbledCookiesInResponse(t *testing.T) {
	// "If a decompressed field section contains multiple cookie field lines,
	// these MUST be concatenated into a single byte string [...]"
	// using the two-byte delimiter of "; "''
	// https://www.rfc-editor.org/rfc/rfc9114.html#section-4.2.1-2
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		req, _ := http.NewRequest("GET", "https://example.tld/", nil)
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)
		st.writeHeaders(http.Header{
			":status": []string{"200"},
			"cookie":  []string{"a=1", "b=2; c=3", "d=4"},
		})
		rt.wantStatus(200)
		rt.wantHeaders(http.Header{
			"Cookie": []string{"a=1; b=2; c=3; d=4"},
		})
	})
}

func TestRoundTripRequestBodySent(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		bodyr, bodyw := io.Pipe()

		req, _ := http.NewRequest("GET", "https://example.tld/", bodyr)
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)

		bodyw.Write([]byte{0, 1, 2, 3, 4})
		st.wantData([]byte{0, 1, 2, 3, 4})

		bodyw.Write([]byte{5, 6, 7})
		st.wantData([]byte{5, 6, 7})

		bodyw.Close()
		st.wantClosed("request body sent")

		st.writeHeaders(http.Header{
			":status": []string{"200"},
		})
		rt.wantStatus(200)
		rt.response().Body.Close()
	})
}

func TestRoundTripRequestBodyErrors(t *testing.T) {
	for _, test := range []struct {
		name          string
		body          io.Reader
		contentLength int64
	}{{
		name:          "too short",
		contentLength: 10,
		body:          bytes.NewReader([]byte{0, 1, 2, 3, 4}),
	}, {
		name:          "too long",
		contentLength: 5,
		body:          bytes.NewReader([]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
	}, {
		name: "read error",
		body: io.MultiReader(
			bytes.NewReader([]byte{0, 1, 2, 3, 4}),
			&testReader{
				readFunc: func([]byte) (int, error) {
					return 0, errors.New("read error")
				},
			},
		),
	}} {
		synctestSubtest(t, test.name, func(t *testing.T) {
			tc := newTestClientConn(t)
			tc.greet()

			req, _ := http.NewRequest("GET", "https://example.tld/", test.body)
			req.ContentLength = test.contentLength
			rt := tc.roundTrip(req)
			st := tc.wantStream(streamTypeRequest)

			// The Transport should send some number of frames before detecting an
			// error in the request body and aborting the request.
			synctest.Wait()
			for {
				_, err := st.readFrameHeader()
				if err != nil {
					var code quic.StreamErrorCode
					if !errors.As(err, &code) {
						t.Fatalf("request stream closed with error %v: want QUIC stream error", err)
					}
					break
				}
				if err := st.discardFrame(); err != nil {
					t.Fatalf("discardFrame: %v", err)
				}
			}

			// RoundTrip returns with an error.
			rt.wantError("request fails due to body error")
		})
	}
}

func TestRoundTripRequestBodyErrorAfterHeaders(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		bodyr, bodyw := io.Pipe()
		req, _ := http.NewRequest("GET", "https://example.tld/", bodyr)
		req.ContentLength = 10
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)

		// Server sends response headers, and RoundTrip returns.
		// The request body hasn't been sent yet.
		st.wantHeaders(nil)
		st.writeHeaders(http.Header{
			":status": []string{"200"},
		})
		rt.wantStatus(200)

		// Write too many bytes to the request body, triggering a request error.
		bodyw.Write(make([]byte, req.ContentLength+1))

		//io.Copy(io.Discard, st)
		st.wantError(quic.StreamErrorCode(errH3InternalError))

		if err := rt.response().Body.Close(); err == nil {
			t.Fatalf("Response.Body.Close() = %v, want error", err)
		}
	})
}

func TestRoundTripExpect100Continue(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		var callCount1xx, callCount100, callCount100Wait int
		trace := &httptrace.ClientTrace{
			Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
				callCount1xx++
				return nil
			},
			Got100Continue: func() {
				callCount100++
			},
			Wait100Continue: func() {
				callCount100Wait++
			},
		}

		tc := newTestClientConn(t)
		tc.greet()
		clientBody := []byte("client's body that will be sent later")
		serverBody := []byte("server's body")

		// Client sends an Expect: 100-continue request.
		req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(t.Context(), trace), "GET", "https://example.tld/", bytes.NewBuffer(clientBody))
		req.Header = http.Header{"Expect": {"100-continue"}}
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)

		// Server reads the header.
		st.wantHeaders(nil)
		st.wantIdle("client has yet to send its body")

		// Server responds with HTTP status 100.
		st.writeHeaders(http.Header{
			":status": []string{"100"},
		})

		// Client sends its body after receiving HTTP status 100 response.
		st.wantData(clientBody)

		// The server sends its response after getting the client's body.
		st.writeHeaders(http.Header{
			":status": []string{"200"},
		})
		st.writeData(serverBody)
		st.stream.stream.CloseWrite()

		// Client receives the response from server.
		rt.wantStatus(200)
		rt.wantBody(serverBody)

		gotCount := []int{callCount1xx, callCount100, callCount100Wait}
		if !slices.Equal(gotCount, []int{1, 1, 1}) {
			t.Errorf("Got1xxResponse, Got100Continue, and Wait100Continue was called %v times respectively, want [1 1 1]", gotCount)
		}
	})
}

func TestRoundTripExpect100ContinueRejected(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		var callCount1xx, callCount100, callCount100Wait int
		trace := &httptrace.ClientTrace{
			Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
				callCount1xx++
				return nil
			},
			Got100Continue: func() {
				callCount100++
			},
			Wait100Continue: func() {
				callCount100Wait++
			},
		}

		tc := newTestClientConn(t)
		tc.greet()

		// Client sends an Expect: 100-continue request.
		req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(t.Context(), trace), "GET", "https://example.tld/", bytes.NewBufferString("client's body"))
		req.Header = http.Header{"Expect": {"100-continue"}}
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)

		// Server reads the header.
		st.wantHeaders(nil)
		st.wantIdle("client has yet to send its body")

		// Server rejects it.
		st.writeHeaders(http.Header{
			":status": []string{"200"},
		})
		st.wantIdle("client does not send its body without getting status 100")
		serverBody := []byte("server's body")
		st.writeData(serverBody)
		st.stream.stream.CloseWrite()

		rt.wantStatus(200)
		rt.wantBody(serverBody)

		gotCount := []int{callCount1xx, callCount100, callCount100Wait}
		if !slices.Equal(gotCount, []int{0, 0, 1}) {
			t.Errorf("Got1xxResponse, Got100Continue, and Wait100Continue was called %v times respectively, want [0 0 1]", gotCount)
		}
	})
}

func TestRoundTripNoBodyClosesStream(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		req, _ := http.NewRequest("PUT", "https://example.tld/", nil)
		tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)

		st.wantHeaders(nil)
		st.wantClosed("no DATA frames to send")
	})
}

func TestRoundTripReadRespWithNoBody(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		// Case 1: we know response body is empty because the server closes the
		// write direction of the stream.
		req, _ := http.NewRequest("GET", "https://example.tld/", nil)
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)
		st.writeHeaders(http.Header{
			":status": {"200"},
		})
		st.stream.stream.CloseWrite()
		rt.wantStatus(200)
		st.wantClosed("request is complete")

		// Case 2: we know response body is empty because the server indicates
		// a Content-Length of 0.
		req, _ = http.NewRequest("GET", "https://example.tld/", nil)
		rt = tc.roundTrip(req)
		st = tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)
		st.writeHeaders(http.Header{
			":status":        {"200"},
			"content-length": {"0"},
		})
		rt.wantStatus(200)
		st.wantClosed("request is complete")

		// Case 3: we know response body is empty because we sent a HEAD
		// request.
		req, _ = http.NewRequest("HEAD", "https://example.tld/", nil)
		rt = tc.roundTrip(req)
		st = tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)
		st.writeHeaders(http.Header{
			":status":        {"200"},
			"content-length": {"1000"},
		})
		rt.wantStatus(200)
		st.wantClosed("request is complete")
	})
}

func TestRoundTripWriteTrailer(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		var req *http.Request
		req, _ = http.NewRequest("POST", "https://example.tld/", io.MultiReader(
			testReader{readFunc: func(_ []byte) (int, error) {
				req.Trailer["Client-Trailer-A"] = []string{"valuea"}
				req.Trailer["Undeclared-Trailer"] = []string{"undeclared"} // Should be ignored.
				return 0, io.EOF
			}},
			strings.NewReader("a body"),
			testReader{readFunc: func(_ []byte) (int, error) {
				req.Trailer["Client-Trailer-B"] = []string{"valueb"}
				req.Trailer["Undeclared-Trailer"] = []string{"undeclared"} // Should be ignored.
				return 0, io.EOF
			}},
		))
		req.Trailer = http.Header{
			"Client-Trailer-A": nil,
			"Client-Trailer-B": nil,
		}
		tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)
		st.wantData([]byte("a body"))
		st.wantHeaders(http.Header{
			"Client-Trailer-A": {"valuea"},
			"Client-Trailer-B": {"valueb"},
		})
		st.wantClosed("request is complete")
	})
}

func TestRoundTripWriteTrailerNoBody(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		var req *http.Request
		req, _ = http.NewRequest("POST", "https://example.tld/", io.MultiReader(
			testReader{readFunc: func(_ []byte) (int, error) {
				req.Trailer["Client-Trailer-A"] = []string{"valuea"}
				req.Trailer["Undeclared-Trailer"] = []string{"undeclared"} // Should be ignored.
				return 0, io.EOF
			}},
			testReader{readFunc: func(_ []byte) (int, error) {
				req.Trailer["Client-Trailer-B"] = []string{"valueb"}
				req.Trailer["Undeclared-Trailer"] = []string{"undeclared"} // Should be ignored.
				return 0, io.EOF
			}},
		))
		req.Trailer = http.Header{
			"Client-Trailer-A": nil,
			"Client-Trailer-B": nil,
		}
		tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)
		st.wantHeaders(nil)
		st.wantHeaders(http.Header{
			"Client-Trailer-A": {"valuea"},
			"Client-Trailer-B": {"valueb"},
		})
		st.wantClosed("request is complete")
	})
}

func TestRoundTripReadTrailer(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		var req *http.Request
		req, _ = http.NewRequest("GET", "https://example.tld/", nil)
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)

		st.wantHeaders(nil)
		st.writeHeaders(http.Header{
			":status": {"200"},
			"trailer": {"Server-Trailer-A, Server-Trailer-B", "server-trailer-c"}, // Should be canonicalized.
		})
		body := []byte("body from server")
		st.writeData(body)
		st.writeHeaders(http.Header{
			"server-trailer-a": {"valuea"},
			// Note that Server-Trailer-B is skipped.
			"server-trailer-c":   {"valuec"},
			"undeclared-trailer": {"undeclared"}, // Should be ignored.
		})

		rt.wantStatus(200)
		// Trailer is stripped off from http.Response.Header and given in http.Response.Trailer.
		rt.wantHeaders(http.Header{})
		rt.wantTrailers(http.Header{
			"Server-Trailer-A": nil,
			"Server-Trailer-B": nil,
			"Server-Trailer-C": nil,
		})

		// Trailer updated after reading the body to EOF.
		rt.wantBody(body)
		rt.wantTrailers(http.Header{
			"Server-Trailer-A": {"valuea"},
			"Server-Trailer-B": nil,
			"Server-Trailer-C": {"valuec"},
		})
		st.wantClosed("request is complete")
	})
}

func TestRoundTripReadTrailerNoBody(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		tc := newTestClientConn(t)
		tc.greet()

		var req *http.Request
		req, _ = http.NewRequest("GET", "https://example.tld/", nil)
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)

		st.wantHeaders(nil)
		st.writeHeaders(http.Header{
			":status":        {"200"},
			"content-length": {"0"},
			"trailer":        {"Server-Trailer-A, Server-Trailer-B", "server-trailer-c"}, // Should be canonicalized.
		})
		st.writeHeaders(http.Header{
			"server-trailer-a": {"valuea"},
			// Note that Server-Trailer-B is skipped.
			"server-trailer-c":   {"valuec"},
			"undeclared-trailer": {"undeclared"}, // Should be ignored.
		})

		rt.wantStatus(200)
		// Trailer is stripped off from http.Response.Header and given in http.Response.Trailer.
		rt.wantHeaders(http.Header{"Content-Length": {"0"}})
		rt.wantTrailers(http.Header{
			"Server-Trailer-A": nil,
			"Server-Trailer-B": nil,
			"Server-Trailer-C": nil,
		})

		// Trailer updated after reading the empty body to EOF.
		rt.wantBody(make([]byte, 0))
		rt.wantTrailers(http.Header{
			"Server-Trailer-A": {"valuea"},
			"Server-Trailer-B": nil,
			"Server-Trailer-C": {"valuec"},
		})
		st.wantClosed("request is complete")
	})
}

func TestRoundTrip103EarlyHints(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		firstHeader := http.Header{
			":status": {"103"},
			"Link":    {"</style.css>; rel=preload; as=style"},
		}
		secondHeader := http.Header{
			":status": {"103"},
			"Link":    {"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script"},
		}

		var respCounter int
		trace := &httptrace.ClientTrace{
			Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
				var wantHeader textproto.MIMEHeader
				switch respCounter {
				case 0:
					wantHeader = textproto.MIMEHeader(firstHeader)
				case 1:
					wantHeader = textproto.MIMEHeader(secondHeader)
				default:
					t.Error("Unexpected 1xx response")
				}
				wantHeader.Del(":status")
				if !reflect.DeepEqual(header, wantHeader) {
					t.Errorf("got %v early hints header, want %v", header, wantHeader)
				}
				respCounter++
				return nil
			},
		}
		req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(t.Context(), trace), "GET", "https://example.tld/", nil)

		tc := newTestClientConn(t)
		tc.greet()
		rt := tc.roundTrip(req)
		st := tc.wantStream(streamTypeRequest)

		st.wantHeaders(nil)
		st.writeHeaders(firstHeader)
		st.writeHeaders(secondHeader)

		st.writeHeaders(http.Header{
			":status": {"200"},
		})
		body := []byte("some body")
		st.writeData(body)
		st.stream.stream.CloseWrite()

		rt.wantStatus(200)
		rt.wantBody(body)
		st.wantClosed("request is complete")
	})
}