Handle multiple requests in single Write()

psiphon/common/transforms/httpTransformer.go

@@ -82,8 +82,7 @@ type HTTPTransformer struct {
 	net.Conn
 }
 
-// Warning: Does not handle chunked encoding and multiple HTTP
-// requests written in a single Write(). Must be called synchronously.
+// Warning: Does not handle chunked encoding. Must be called synchronously.
 func (t *HTTPTransformer) Write(b []byte) (int, error) {
 
 	if t.state == httpTransformerReadHeader {
@@ -170,14 +169,29 @@ func (t *HTTPTransformer) Write(b []byte) (int, error) {
 			err = t.writeBuffer()
 
 			if t.remain > 0 {
+				t.state = httpTransformerReadWriteBody
+			} else {
 				// Entire request, header and body, has been written. Return to
 				// waiting for next HTTP request header to arrive.
-				t.state = httpTransformerReadWriteBody
+				if len(t.b) > 0 {
+					// Return the number of bytes written to the underlying
+					// Conn and clear t.b instead of calling t.Write() with any
+					// remaining bytes of t.b. The caller must call Write()
+					// again with the unwritten, and unbuffered, bytes of b.
+					// Since t.remain = 0 it is guaranteed that
+					// len(b) - len(t.b) >= 0 because len(t.b) is the number of
+					// subsequent request bytes and len(b) is the number of
+					// trailing bytes of the current request plus the
+					// subsequent request bytes.
+					written := len(b) - len(t.b)
+					t.b = nil
+					return written, err
+				}
 			}
 
 			if err != nil {
 				// b buffered in t.b
-				return len(b), errors.Trace(err)
+				return len(b), err
 			}
 		}
 
@@ -196,21 +210,24 @@ func (t *HTTPTransformer) Write(b []byte) (int, error) {
 		return 0, errors.Trace(err)
 	}
 
-	n, err := t.Conn.Write(b)
-
-	if uint64(n) > t.remain {
-		// Attempt to recover by resetting t.remain. If b contains bytes of
-		// subsequent request(s) (should never happen), then these request(s)
-		// may be mangled.
-		t.remain = 0
-		return n, errors.TraceNew("t.remain - uint64(n) underflows")
+	bytesToWrite := uint64(len(b))
+	if bytesToWrite > t.remain {
+		bytesToWrite = t.remain
 	}
 
+	n, err := t.Conn.Write(b[:bytesToWrite])
+
+	// Do not need to check for underflow because n <= t.remain
 	t.remain -= uint64(n)
 
 	if t.remain <= 0 {
 		// Entire request, header and body, has been written. Return to
 		// waiting for next HTTP request header to arrive.
+		//
+		// Return the number of bytes written to the underlying Conn instead of
+		// calling t.Write() with any remaining bytes of b which were not
+		// written or buffered, i.e. when n < len(b). The caller must call
+		// Write() again with the unwritten, and unbuffered, bytes of b.
 		t.state = httpTransformerReadHeader
 		t.remain = 0
 	}
@@ -219,17 +236,16 @@ func (t *HTTPTransformer) Write(b []byte) (int, error) {
 }
 
 func (t *HTTPTransformer) writeBuffer() error {
-	for len(t.b) > 0 {
-		n, err := t.Conn.Write(t.b)
-
-		if uint64(n) > t.remain {
-			// Attempt to recover by resetting t.remain. If t.b contains bytes
-			// of subsequent request(s) (should never happen), then these
-			// request(s) may be mangled.
-			t.remain = 0
-			return errors.TraceNew("t.remain - uint64(n) underflows")
+	for len(t.b) > 0 && t.remain > 0 {
+
+		bytesToWrite := uint64(len(t.b))
+		if bytesToWrite > t.remain {
+			bytesToWrite = t.remain
 		}
 
+		n, err := t.Conn.Write(t.b[:bytesToWrite])
+
+		// Do not need to check for underflow because n <= t.remain
 		t.remain -= uint64(n)
 
 		if n == len(t.b) {

psiphon/common/transforms/httpTransformer_test.go

@@ -125,24 +125,39 @@ func TestHTTPTransformerHTTPRequest(t *testing.T) {
 			connWriteLimit: 1,
 			connWriteErrs:  []error{errors.New("err1"), errors.New("err2"), errors.New("err3")},
 		},
-		// Multiple HTTP requests written in a single write not supported so an
-		// error is expected.
+		// Multiple HTTP requests written in a single write.
 		{
 			name:       "multiple HTTP requests written in a single write",
 			input:      "HTTP 1.1\r\nContent-Length: 4\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 2\r\n\r\n12",
 			wantOutput: "HTTP 1.1\r\nContent-Length: 4\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 2\r\n\r\n12",
 			chunkSize:  999,
-			wantError:  errors.New("t.remain - uint64(n) underflows"),
 		},
-		// Multiple HTTP requests written in a single write not supported so an
-		// error is expected because a write will occur where it contains both
-		// the end of the previous HTTP request and the start of a new one.
+		// Multiple HTTP requests written in a single write. A write will occur
+		// where it contains both the end of the previous HTTP request and the
+		// start of a new one.
 		{
 			name:       "multiple HTTP requests written in chunks",
 			input:      "HTTP 1.1\r\nContent-Length: 4\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 2\r\n\r\n12",
 			wantOutput: "HTTP 1.1\r\nContent-Length: 4\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 2\r\n\r\n12",
 			chunkSize:  3,
-			wantError:  errors.New("t.remain - uint64(n) underflows"),
+		},
+		// Multiple HTTP requests written in a single write with transform.
+		{
+			name:       "multiple HTTP requests written in a single write",
+			input:      "HTTP 1.1\r\nContent-Length: 4\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 4\r\n\r\n12HTTP 1.1\r\nContent-Length: 4\r\n\r\n34",
+			wantOutput: "HTTP 1.1\r\nContent-Length: 100\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 100\r\n\r\n12HTTP 1.1\r\nContent-Length: 100\r\n\r\n34",
+			chunkSize:  999,
+			transform:  Spec{[2]string{"4", "100"}},
+		},
+		// Multiple HTTP requests written in a single write with transform. A
+		// write will occur where it contains both the end of the previous HTTP
+		// request and the start of a new one.
+		{
+			name:       "multiple HTTP requests written in chunks",
+			input:      "HTTP 1.1\r\nContent-Length: 4\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 4\r\n\r\n12",
+			wantOutput: "HTTP 1.1\r\nContent-Length: 100\r\n\r\nabcdHTTP 1.1\r\nContent-Length: 100\r\n\r\n12",
+			chunkSize:  3,
+			transform:  Spec{[2]string{"4", "100"}},
 		},
 	}