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 = .{},
2121/// it will first rescan the system for root certificates.
2222next_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
2435pub const Connection = struct {
2536 stream: net.Stream,
2637 /// 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.
2839 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
3046 pub const Protocol = enum { plain, tls };
3147
......@@ -56,6 +72,17 @@ pub const Connection = struct {
5672 .tls => return conn.tls_client.write(conn.stream, buffer),
5773 }
5874 }
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 }
5986};
6087
6188/// TODO: emit error.UnexpectedEndOfStream or something like that when the read
......@@ -63,7 +90,7 @@ pub const Connection = struct {
6390/// close_notify protection on underlying TLS streams.
6491pub const Request = struct {
6592 client: *Client,
66 connection: Connection,
93 connection: *ConnectionNode,
6794 redirects_left: u32,
6895 response: Response,
6996 /// These are stored in Request so that they are available when following
......@@ -79,6 +106,7 @@ pub const Request = struct {
79106 header_bytes: std.ArrayListUnmanaged(u8),
80107 max_header_bytes: usize,
81108 next_chunk_length: u64,
109 done: bool,
82110
83111 pub const Headers = struct {
84112 status: http.Status,
......@@ -86,6 +114,7 @@ pub const Request = struct {
86114 location: ?[]const u8 = null,
87115 content_length: ?u64 = null,
88116 transfer_encoding: ?http.TransferEncoding = null,
117 connection_close: bool = true,
89118
90119 pub fn parse(bytes: []const u8) !Response.Headers {
91120 var it = mem.split(u8, bytes[0 .. bytes.len - 4], "\r\n");
......@@ -126,6 +155,14 @@ pub const Request = struct {
126155 if (headers.transfer_encoding != null) return error.HttpHeadersInvalid;
127156 headers.transfer_encoding = std.meta.stringToEnum(http.TransferEncoding, header_value) orelse
128157 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 }
129166 }
130167 }
131168
......@@ -185,10 +222,10 @@ pub const Request = struct {
185222 chunk_r,
186223 chunk_data,
187224
188 pub fn zeroMeansEnd(state: State) bool {
189 return switch (state) {
190 .finished, .chunk_data => true,
191 else => false,
225 pub fn isContent(self: State) bool {
226 return switch (self) {
227 .invalid, .start, .seen_r, .seen_rn, .seen_rnr => false,
228 .finished, .chunk_size_prefix_r, .chunk_size_prefix_n, .chunk_size, .chunk_r, .chunk_data => true,
192229 };
193230 }
194231 };
......@@ -201,6 +238,7 @@ pub const Request = struct {
201238 .max_header_bytes = max,
202239 .header_bytes_owned = true,
203240 .next_chunk_length = undefined,
241 .done = false,
204242 };
205243 }
206244
......@@ -212,6 +250,7 @@ pub const Request = struct {
212250 .max_header_bytes = buf.len,
213251 .header_bytes_owned = false,
214252 .next_chunk_length = undefined,
253 .done = false,
215254 };
216255 }
217256
......@@ -501,6 +540,7 @@ pub const Request = struct {
501540 pub const Headers = struct {
502541 version: http.Version = .@"HTTP/1.1",
503542 method: http.Method = .GET,
543 connection_close: bool = false,
504544 };
505545
506546 pub const Options = struct {
......@@ -545,6 +585,7 @@ pub const Request = struct {
545585 HttpHeadersExceededSizeLimit,
546586 HttpRedirectMissingLocation,
547587 HttpTransferEncodingUnsupported,
588 HttpConnectionHeaderUnsupported,
548589 HttpContentLengthUnknown,
549590 TooManyHttpRedirects,
550591 ShortHttpStatusLine,
......@@ -669,8 +710,9 @@ pub const Request = struct {
669710 assert(len <= buffer.len);
670711 var index: usize = 0;
671712 while (index < len) {
672 const zero_means_end = req.response.state.zeroMeansEnd();
673713 const amt = try readAdvanced(req, buffer[index..]);
714 const zero_means_end = req.response.done and req.response.headers.status.class() != .redirect;
715
674716 if (amt == 0 and zero_means_end) break;
675717 index += amt;
676718 }
......@@ -680,7 +722,29 @@ pub const Request = struct {
680722 /// This one can return 0 without meaning EOF.
681723 /// TODO change to readvAdvanced
682724 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)];
684748 var out_index: usize = 0;
685749 while (true) {
686750 switch (req.response.state) {
......@@ -698,24 +762,10 @@ pub const Request = struct {
698762 if (req.response.state == .finished) {
699763 req.response.headers = try Response.Headers.parse(req.response.header_bytes.items);
700764
701 if (req.response.headers.status.class() == .redirect) {
702 if (req.redirects_left == 0) return error.TooManyHttpRedirects;
703 const location = req.response.headers.location orelse
704 return error.HttpRedirectMissingLocation;
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;
765 if (req.response.headers.connection_close == true) {
766 req.connection.data.unusable = true;
767 } else {
768 req.connection.data.unusable = false;
719769 }
720770
721771 if (req.response.headers.transfer_encoding) |transfer_encoding| {
......@@ -742,11 +792,29 @@ pub const Request = struct {
742792 return 0;
743793 },
744794 .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
745813 if (in.ptr == buffer.ptr) {
746 return in.len;
814 return sub_amt;
747815 } else {
748 mem.copy(u8, buffer[out_index..], in);
749 return out_index + in.len;
816 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);
817 return out_index + sub_amt;
750818 }
751819 },
752820 .chunk_size_prefix_r => switch (in.len) {
......@@ -793,7 +861,10 @@ pub const Request = struct {
793861 .invalid => return error.HttpHeadersInvalid,
794862 .chunk_data => {
795863 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
797868 return out_index;
798869 }
799870 in = in[i..];
......@@ -807,20 +878,27 @@ pub const Request = struct {
807878 // TODO https://github.com/ziglang/zig/issues/14039
808879 const sub_amt = @intCast(usize, @min(req.response.next_chunk_length, in.len));
809880 req.response.next_chunk_length -= sub_amt;
810 if (req.response.next_chunk_length > 0) {
811 if (in.ptr == buffer.ptr) {
812 return sub_amt;
813 } else {
814 mem.copy(u8, buffer[out_index..], in[0..sub_amt]);
815 out_index += sub_amt;
816 return out_index;
817 }
881
882 if (req.response.next_chunk_length == 0) {
883 req.response.state = .chunk_size_prefix_r;
884 in = in[sub_amt..];
885
886 if (req.response.headers.status.class() == .redirect) continue;
887
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;
818901 }
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;
824902 },
825903 }
826904 }
......@@ -844,24 +922,52 @@ pub const Request = struct {
844922};
845923
846924pub 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
847934 client.ca_bundle.deinit(client.allocator);
848935 client.* = undefined;
849936}
850937
851pub fn connect(client: *Client, host: []const u8, port: u16, protocol: Connection.Protocol) !Connection {
852 var conn: Connection = .{
938pub fn connect(client: *Client, host: []const u8, port: u16, protocol: Connection.Protocol) !*ConnectionNode {
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 = .{
853957 .stream = try net.tcpConnectToHost(client.allocator, host, port),
854958 .tls_client = undefined,
855959 .protocol = protocol,
856 };
960 .host = try client.allocator.dupe(u8, host),
961 .port = port,
962 } };
857963
858964 switch (protocol) {
859965 .plain => {},
860966 .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);
862968 // This is appropriate for HTTPS because the HTTP headers contain
863969 // 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;
865971 },
866972 }
867973
......@@ -908,10 +1014,15 @@ pub fn request(client: *Client, uri: Uri, headers: Request.Headers, options: Req
9081014 try h.appendSlice(@tagName(headers.version));
9091015 try h.appendSlice("\r\nHost: ");
9101016 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
9131024 const header_bytes = h.slice();
914 try req.connection.writeAll(header_bytes);
1025 try req.connection.data.writeAll(header_bytes);
9151026 }
9161027
9171028 return req;