authorgravatar for truemedian@gmail.comNameless <truemedian@gmail.com> 2023-03-02 12:45:34-06:00
committergravatar for truemedian@gmail.comNameless <truemedian@gmail.com> 2023-03-09 14:54:23-06:00
logafb26f4e6b39431001eff75cc8ce19144cb5301a
treea09b945509968e8211a4a7be16f62157f0c0a303
parent95f6a5935a675efe6d30bc2388e7a0bc6b742c6d
signature Commit is signed but in an unrecognized format.

std.http: add connection pooling and make keep-alive requests by default

1 files changed, 161 insertions(+), 50 deletions(-)

lib/std/http/Client.zig+161-50
...@@ -21,11 +21,27 @@ ca_bundle: std.crypto.Certificate.Bundle = .{},...@@ -21,11 +21,27 @@ ca_bundle: std.crypto.Certificate.Bundle = .{},
21/// it will first rescan the system for root certificates.21/// it will first rescan the system for root certificates.
22next_https_rescan_certs: bool = true,22next_https_rescan_certs: bool = true,
2323
24connection_pool: std.TailQueue(Connection) = .{},
25
26const ConnectionPool = std.TailQueue(Connection);
27const ConnectionNode = ConnectionPool.Node;
28
29pub fn release(client: *Client, node: *ConnectionNode) void {
30 if (node.data.unusable) return node.data.close(client);
31
32 client.connection_pool.append(node);
33}
34
24pub const Connection = struct {35pub const Connection = struct {
25 stream: net.Stream,36 stream: net.Stream,
26 /// undefined unless protocol is tls.37 /// undefined unless protocol is tls.
27 tls_client: std.crypto.tls.Client,38 tls_client: std.crypto.tls.Client, // TODO: allocate this, it's currently 16 KB.
28 protocol: Protocol,39 protocol: Protocol,
40 host: []u8,
41 port: u16,
42
43 // This connection has been part of a non keepalive request and cannot be added to the pool.
44 unusable: bool = false,
2945
30 pub const Protocol = enum { plain, tls };46 pub const Protocol = enum { plain, tls };
3147
...@@ -56,6 +72,17 @@ pub const Connection = struct {...@@ -56,6 +72,17 @@ pub const Connection = struct {
56 .tls => return conn.tls_client.write(conn.stream, buffer),72 .tls => return conn.tls_client.write(conn.stream, buffer),
57 }73 }
58 }74 }
75
76 pub fn close(conn: *Connection, client: *const Client) void {
77 if (conn.protocol == .tls) {
78 // try to cleanly close the TLS connection, for any server that cares.
79 _ = conn.tls_client.writeEnd(conn.stream, "", true) catch {};
80 }
81
82 conn.stream.close();
83
84 client.allocator.free(conn.host);
85 }
59};86};
6087
61/// TODO: emit error.UnexpectedEndOfStream or something like that when the read88/// TODO: emit error.UnexpectedEndOfStream or something like that when the read
...@@ -63,7 +90,7 @@ pub const Connection = struct {...@@ -63,7 +90,7 @@ pub const Connection = struct {
63/// close_notify protection on underlying TLS streams.90/// close_notify protection on underlying TLS streams.
64pub const Request = struct {91pub const Request = struct {
65 client: *Client,92 client: *Client,
66 connection: Connection,93 connection: *ConnectionNode,
67 redirects_left: u32,94 redirects_left: u32,
68 response: Response,95 response: Response,
69 /// These are stored in Request so that they are available when following96 /// These are stored in Request so that they are available when following
...@@ -79,6 +106,7 @@ pub const Request = struct {...@@ -79,6 +106,7 @@ pub const Request = struct {
79 header_bytes: std.ArrayListUnmanaged(u8),106 header_bytes: std.ArrayListUnmanaged(u8),
80 max_header_bytes: usize,107 max_header_bytes: usize,
81 next_chunk_length: u64,108 next_chunk_length: u64,
109 done: bool,
82110
83 pub const Headers = struct {111 pub const Headers = struct {
84 status: http.Status,112 status: http.Status,
...@@ -86,6 +114,7 @@ pub const Request = struct {...@@ -86,6 +114,7 @@ pub const Request = struct {
86 location: ?[]const u8 = null,114 location: ?[]const u8 = null,
87 content_length: ?u64 = null,115 content_length: ?u64 = null,
88 transfer_encoding: ?http.TransferEncoding = null,116 transfer_encoding: ?http.TransferEncoding = null,
117 connection_close: bool = true,
89118
90 pub fn parse(bytes: []const u8) !Response.Headers {119 pub fn parse(bytes: []const u8) !Response.Headers {
91 var it = mem.split(u8, bytes[0 .. bytes.len - 4], "\r\n");120 var it = mem.split(u8, bytes[0 .. bytes.len - 4], "\r\n");
...@@ -126,6 +155,14 @@ pub const Request = struct {...@@ -126,6 +155,14 @@ pub const Request = struct {
126 if (headers.transfer_encoding != null) return error.HttpHeadersInvalid;155 if (headers.transfer_encoding != null) return error.HttpHeadersInvalid;
127 headers.transfer_encoding = std.meta.stringToEnum(http.TransferEncoding, header_value) orelse156 headers.transfer_encoding = std.meta.stringToEnum(http.TransferEncoding, header_value) orelse
128 return error.HttpTransferEncodingUnsupported;157 return error.HttpTransferEncodingUnsupported;
158 } else if (std.ascii.eqlIgnoreCase(header_name, "connection")) {
159 if (std.ascii.eqlIgnoreCase(header_value, "keep-alive")) {
160 headers.connection_close = false;
161 } else if (std.ascii.eqlIgnoreCase(header_value, "close")) {
162 headers.connection_close = true;
163 } else {
164 return error.HttpConnectionHeaderUnsupported;
165 }
129 }166 }
130 }167 }
131168
...@@ -185,10 +222,10 @@ pub const Request = struct {...@@ -185,10 +222,10 @@ pub const Request = struct {
185 chunk_r,222 chunk_r,
186 chunk_data,223 chunk_data,
187224
188 pub fn zeroMeansEnd(state: State) bool {225 pub fn isContent(self: State) bool {
189 return switch (state) {226 return switch (self) {
190 .finished, .chunk_data => true,227 .invalid, .start, .seen_r, .seen_rn, .seen_rnr => false,
191 else => false,228 .finished, .chunk_size_prefix_r, .chunk_size_prefix_n, .chunk_size, .chunk_r, .chunk_data => true,
192 };229 };
193 }230 }
194 };231 };
...@@ -201,6 +238,7 @@ pub const Request = struct {...@@ -201,6 +238,7 @@ pub const Request = struct {
201 .max_header_bytes = max,238 .max_header_bytes = max,
202 .header_bytes_owned = true,239 .header_bytes_owned = true,
203 .next_chunk_length = undefined,240 .next_chunk_length = undefined,
241 .done = false,
204 };242 };
205 }243 }
206244
...@@ -212,6 +250,7 @@ pub const Request = struct {...@@ -212,6 +250,7 @@ pub const Request = struct {
212 .max_header_bytes = buf.len,250 .max_header_bytes = buf.len,
213 .header_bytes_owned = false,251 .header_bytes_owned = false,
214 .next_chunk_length = undefined,252 .next_chunk_length = undefined,
253 .done = false,
215 };254 };
216 }255 }
217256
...@@ -501,6 +540,7 @@ pub const Request = struct {...@@ -501,6 +540,7 @@ pub const Request = struct {
501 pub const Headers = struct {540 pub const Headers = struct {
502 version: http.Version = .@"HTTP/1.1",541 version: http.Version = .@"HTTP/1.1",
503 method: http.Method = .GET,542 method: http.Method = .GET,
543 connection_close: bool = false,
504 };544 };
505545
506 pub const Options = struct {546 pub const Options = struct {
...@@ -545,6 +585,7 @@ pub const Request = struct {...@@ -545,6 +585,7 @@ pub const Request = struct {
545 HttpHeadersExceededSizeLimit,585 HttpHeadersExceededSizeLimit,
546 HttpRedirectMissingLocation,586 HttpRedirectMissingLocation,
547 HttpTransferEncodingUnsupported,587 HttpTransferEncodingUnsupported,
588 HttpConnectionHeaderUnsupported,
548 HttpContentLengthUnknown,589 HttpContentLengthUnknown,
549 TooManyHttpRedirects,590 TooManyHttpRedirects,
550 ShortHttpStatusLine,591 ShortHttpStatusLine,
...@@ -669,8 +710,9 @@ pub const Request = struct {...@@ -669,8 +710,9 @@ pub const Request = struct {
669 assert(len <= buffer.len);710 assert(len <= buffer.len);
670 var index: usize = 0;711 var index: usize = 0;
671 while (index < len) {712 while (index < len) {
672 const zero_means_end = req.response.state.zeroMeansEnd();
673 const amt = try readAdvanced(req, buffer[index..]);713 const amt = try readAdvanced(req, buffer[index..]);
714 const zero_means_end = req.response.done and req.response.headers.status.class() != .redirect;
715
674 if (amt == 0 and zero_means_end) break;716 if (amt == 0 and zero_means_end) break;
675 index += amt;717 index += amt;
676 }718 }
...@@ -680,7 +722,29 @@ pub const Request = struct {...@@ -680,7 +722,29 @@ pub const Request = struct {
680 /// This one can return 0 without meaning EOF.722 /// This one can return 0 without meaning EOF.
681 /// TODO change to readvAdvanced723 /// TODO change to readvAdvanced
682 pub fn readAdvanced(req: *Request, buffer: []u8) !usize {724 pub fn readAdvanced(req: *Request, buffer: []u8) !usize {
683 var in = buffer[0..try req.connection.read(buffer)];725 if (req.response.done) {
726 if (req.response.headers.status.class() == .redirect) {
727 if (req.redirects_left == 0) return error.TooManyHttpRedirects;
728
729 const location = req.response.headers.location orelse
730 return error.HttpRedirectMissingLocation;
731 const new_url = try std.Uri.parse(location);
732 const new_req = try req.client.request(new_url, req.headers, .{
733 .max_redirects = req.redirects_left - 1,
734 .header_strategy = if (req.response.header_bytes_owned) .{
735 .dynamic = req.response.max_header_bytes,
736 } else .{
737 .static = req.response.header_bytes.unusedCapacitySlice(),
738 },
739 });
740 req.deinit();
741 req.* = new_req;
742 } else {
743 return 0;
744 }
745 }
746
747 var in = buffer[0..try req.connection.data.read(buffer)];
684 var out_index: usize = 0;748 var out_index: usize = 0;
685 while (true) {749 while (true) {
686 switch (req.response.state) {750 switch (req.response.state) {
...@@ -698,24 +762,10 @@ pub const Request = struct {...@@ -698,24 +762,10 @@ pub const Request = struct {
698 if (req.response.state == .finished) {762 if (req.response.state == .finished) {
699 req.response.headers = try Response.Headers.parse(req.response.header_bytes.items);763 req.response.headers = try Response.Headers.parse(req.response.header_bytes.items);
700764
701 if (req.response.headers.status.class() == .redirect) {765 if (req.response.headers.connection_close == true) {
702 if (req.redirects_left == 0) return error.TooManyHttpRedirects;766 req.connection.data.unusable = true;
703 const location = req.response.headers.location orelse767 } else {
704 return error.HttpRedirectMissingLocation;768 req.connection.data.unusable = false;
705 const new_url = try std.Uri.parse(location);
706 const new_req = try req.client.request(new_url, req.headers, .{
707 .max_redirects = req.redirects_left - 1,
708 .header_strategy = if (req.response.header_bytes_owned) .{
709 .dynamic = req.response.max_header_bytes,
710 } else .{
711 .static = req.response.header_bytes.unusedCapacitySlice(),
712 },
713 });
714 req.deinit();
715 req.* = new_req;
716 assert(out_index == 0);
717 in = buffer[0..try req.connection.read(buffer)];
718 continue;
719 }769 }
720770
721 if (req.response.headers.transfer_encoding) |transfer_encoding| {771 if (req.response.headers.transfer_encoding) |transfer_encoding| {
...@@ -742,11 +792,29 @@ pub const Request = struct {...@@ -742,11 +792,29 @@ pub const Request = struct {
742 return 0;792 return 0;
743 },793 },
744 .finished => {794 .finished => {
795 const sub_amt = @intCast(usize, @min(req.response.next_chunk_length, in.len));
796 req.response.next_chunk_length -= sub_amt;
797
798 if (req.response.next_chunk_length == 0) {
799 req.client.release(req.connection);
800 req.connection = undefined;
801
802 req.response.done = true;
803 assert(in.len == sub_amt); // TODO: figure out how to not read more than necessary.
804
805 if (req.response.state.isContent() and req.response.headers.status.class() == .redirect) return 0;
806
807 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);
808 return out_index + sub_amt;
809 }
810
811 if (req.response.state.isContent() and req.response.headers.status.class() == .redirect) return 0;
812
745 if (in.ptr == buffer.ptr) {813 if (in.ptr == buffer.ptr) {
746 return in.len;814 return sub_amt;
747 } else {815 } else {
748 mem.copy(u8, buffer[out_index..], in);816 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);
749 return out_index + in.len;817 return out_index + sub_amt;
750 }818 }
751 },819 },
752 .chunk_size_prefix_r => switch (in.len) {820 .chunk_size_prefix_r => switch (in.len) {
...@@ -793,7 +861,10 @@ pub const Request = struct {...@@ -793,7 +861,10 @@ pub const Request = struct {
793 .invalid => return error.HttpHeadersInvalid,861 .invalid => return error.HttpHeadersInvalid,
794 .chunk_data => {862 .chunk_data => {
795 if (req.response.next_chunk_length == 0) {863 if (req.response.next_chunk_length == 0) {
796 req.response.state = .start;864 req.response.done = true;
865 req.client.release(req.connection);
866 req.connection = undefined;
867
797 return out_index;868 return out_index;
798 }869 }
799 in = in[i..];870 in = in[i..];
...@@ -807,20 +878,27 @@ pub const Request = struct {...@@ -807,20 +878,27 @@ pub const Request = struct {
807 // TODO https://github.com/ziglang/zig/issues/14039878 // TODO https://github.com/ziglang/zig/issues/14039
808 const sub_amt = @intCast(usize, @min(req.response.next_chunk_length, in.len));879 const sub_amt = @intCast(usize, @min(req.response.next_chunk_length, in.len));
809 req.response.next_chunk_length -= sub_amt;880 req.response.next_chunk_length -= sub_amt;
810 if (req.response.next_chunk_length > 0) {881
811 if (in.ptr == buffer.ptr) {882 if (req.response.next_chunk_length == 0) {
812 return sub_amt;883 req.response.state = .chunk_size_prefix_r;
813 } else {884 in = in[sub_amt..];
814 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);885
815 out_index += sub_amt;886 if (req.response.headers.status.class() == .redirect) continue;
816 return out_index;887
817 }888 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);
889 out_index += sub_amt;
890 continue;
891 }
892
893 if (req.response.headers.status.class() == .redirect) return 0;
894
895 if (in.ptr == buffer.ptr) {
896 return sub_amt;
897 } else {
898 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);
899 out_index += sub_amt;
900 return out_index;
818 }901 }
819 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);
820 out_index += sub_amt;
821 req.response.state = .chunk_size_prefix_r;
822 in = in[sub_amt..];
823 continue;
824 },902 },
825 }903 }
826 }904 }
...@@ -844,24 +922,52 @@ pub const Request = struct {...@@ -844,24 +922,52 @@ pub const Request = struct {
844};922};
845923
846pub fn deinit(client: *Client) void {924pub fn deinit(client: *Client) void {
925 var next = client.connection_pool.first;
926 while (next) |node| {
927 next = node.next;
928
929 node.data.close(client);
930
931 client.allocator.destroy(node);
932 }
933
847 client.ca_bundle.deinit(client.allocator);934 client.ca_bundle.deinit(client.allocator);
848 client.* = undefined;935 client.* = undefined;
849}936}
850937
851pub fn connect(client: *Client, host: []const u8, port: u16, protocol: Connection.Protocol) !Connection {938pub fn connect(client: *Client, host: []const u8, port: u16, protocol: Connection.Protocol) !*ConnectionNode {
852 var conn: Connection = .{939 var potential = client.connection_pool.last;
940 while (potential) |node| {
941 const same_host = mem.eql(u8, node.data.host, host);
942 const same_port = node.data.port == port;
943 const same_protocol = node.data.protocol == protocol;
944
945 if (same_host and same_port and same_protocol) {
946 client.connection_pool.remove(node);
947 return node;
948 }
949
950 potential = node.prev;
951 }
952
953 const conn = try client.allocator.create(ConnectionNode);
954 errdefer client.allocator.destroy(conn);
955
956 conn.* = .{ .data = .{
853 .stream = try net.tcpConnectToHost(client.allocator, host, port),957 .stream = try net.tcpConnectToHost(client.allocator, host, port),
854 .tls_client = undefined,958 .tls_client = undefined,
855 .protocol = protocol,959 .protocol = protocol,
856 };960 .host = try client.allocator.dupe(u8, host),
961 .port = port,
962 } };
857963
858 switch (protocol) {964 switch (protocol) {
859 .plain => {},965 .plain => {},
860 .tls => {966 .tls => {
861 conn.tls_client = try std.crypto.tls.Client.init(conn.stream, client.ca_bundle, host);967 conn.data.tls_client = try std.crypto.tls.Client.init(conn.data.stream, client.ca_bundle, host);
862 // This is appropriate for HTTPS because the HTTP headers contain968 // This is appropriate for HTTPS because the HTTP headers contain
863 // the content length which is used to detect truncation attacks.969 // the content length which is used to detect truncation attacks.
864 conn.tls_client.allow_truncation_attacks = true;970 conn.data.tls_client.allow_truncation_attacks = true;
865 },971 },
866 }972 }
867973
...@@ -908,10 +1014,15 @@ pub fn request(client: *Client, uri: Uri, headers: Request.Headers, options: Req...@@ -908,10 +1014,15 @@ pub fn request(client: *Client, uri: Uri, headers: Request.Headers, options: Req
908 try h.appendSlice(@tagName(headers.version));1014 try h.appendSlice(@tagName(headers.version));
909 try h.appendSlice("\r\nHost: ");1015 try h.appendSlice("\r\nHost: ");
910 try h.appendSlice(host);1016 try h.appendSlice(host);
911 try h.appendSlice("\r\nConnection: close\r\n\r\n");1017 if (headers.connection_close) {
1018 try h.appendSlice("\r\nConnection: close");
1019 } else {
1020 try h.appendSlice("\r\nConnection: keep-alive");
1021 }
1022 try h.appendSlice("\r\n\r\n");
9121023
913 const header_bytes = h.slice();1024 const header_bytes = h.slice();
914 try req.connection.writeAll(header_bytes);1025 try req.connection.data.writeAll(header_bytes);
915 }1026 }
9161027
917 return req;1028 return req;