1 files changed, 28 insertions, 26 deletions

diff --git a/src/main.zig b/src/main.zig
index bce3ca6..4f0f964 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -64,25 +64,21 @@ const SaprusMessage = union(SaprusPacketType) {
 
         switch (self) {
             .relay => |r| {
-                var b64_payload_buf = std.ArrayList(u8).init(allocator);
-                defer b64_payload_buf.deinit();
-                const buf_w = b64_payload_buf.writer();
-                try base64Enc.encodeWriter(buf_w, r.payload);
+                const payload_list = try encodeToList(r.payload, allocator);
+                defer payload_list.deinit();
 
                 try w.writeStructEndian(r.header, .big);
-                try w.writeInt(u16, @intCast(b64_payload_buf.items.len), .big);
-                try w.writeAll(b64_payload_buf.items);
+                try w.writeInt(u16, @intCast(payload_list.items.len), .big);
+                try w.writeAll(payload_list.items);
             },
             .file_transfer => return SaprusError.NotImplementedSaprusType,
             .connection => |c| {
-                var b64_payload_buf = std.ArrayList(u8).init(allocator);
-                defer b64_payload_buf.deinit();
-                const buf_w = b64_payload_buf.writer();
-                try base64Enc.encodeWriter(buf_w, c.payload);
+                const payload_list = try encodeToList(c.payload, allocator);
+                defer payload_list.deinit();
 
                 try w.writeStructEndian(c.header, .big);
-                try w.writeInt(u16, @intCast(b64_payload_buf.items.len), .big);
-                try w.writeAll(b64_payload_buf.items);
+                try w.writeInt(u16, @intCast(payload_list.items.len), .big);
+                try w.writeAll(payload_list.items);
             },
         }
 
@@ -97,13 +93,7 @@ const SaprusMessage = union(SaprusPacketType) {
             .relay => {
                 const header = try r.readStructEndian(Relay.Header, .big);
                 const len = try r.readInt(u16, .big);
-
-                var payload_b64_buf = std.ArrayList(u8).init(allocator);
-                defer payload_b64_buf.deinit();
-                try r.readAllArrayList(&payload_b64_buf, len);
-
-                const payload = try allocator.alloc(u8, try base64Dec.calcSizeForSlice(payload_b64_buf.items));
-                try base64Dec.decode(payload, payload_b64_buf.items);
+                const payload = try decodeFromReader(r, len, allocator);
 
                 return .{
                     .relay = .{
@@ -116,13 +106,7 @@ const SaprusMessage = union(SaprusPacketType) {
             .connection => {
                 const header = try r.readStructEndian(Connection.Header, .big);
                 const len = try r.readInt(u16, .big);
-
-                var payload_b64_buf = std.ArrayList(u8).init(allocator);
-                defer payload_b64_buf.deinit();
-                try r.readAllArrayList(&payload_b64_buf, len);
-
-                const payload = try allocator.alloc(u8, try base64Dec.calcSizeForSlice(payload_b64_buf.items));
-                try base64Dec.decode(payload, payload_b64_buf.items);
+                const payload = try decodeFromReader(r, len, allocator);
 
                 return .{
                     .connection = .{
@@ -178,10 +162,28 @@ pub fn main() !void {
     _ = try sock.sendTo(dest_addr, msg_bytes);
 }
 
+fn encodeToList(buf: []const u8, allocator: Allocator) !std.ArrayList(u8) {
+    var b64_buf = std.ArrayList(u8).init(allocator);
+    const buf_w = b64_buf.writer();
+    try base64Enc.encodeWriter(buf_w, buf);
+    return b64_buf;
+}
+
+fn decodeFromReader(reader: StringReader, len: u16, allocator: Allocator) ![]const u8 {
+    var b64_buf = std.ArrayList(u8).init(allocator);
+    defer b64_buf.deinit();
+    try reader.readAllArrayList(&b64_buf, len);
+
+    const payload = try allocator.alloc(u8, try base64Dec.calcSizeForSlice(b64_buf.items));
+    try base64Dec.decode(payload, b64_buf.items);
+    return payload;
+}
+
 const builtin = @import("builtin");
 const std = @import("std");
 const Allocator = std.mem.Allocator;
 const DebugAllocator = std.heap.DebugAllocator(.{});
+const StringReader = std.io.FixedBufferStream([]const u8).Reader;
 
 const network = @import("network");