Use entropy cache for websocket mask

2 files changed, 17 insertions, 60 deletions

diff --git a/src/bun.js/rare_data.zig b/src/bun.js/rare_data.zig
index 8237d99a9..8686fa1b4 100644
--- a/src/bun.js/rare_data.zig
+++ b/src/bun.js/rare_data.zig
@@ -15,9 +15,8 @@ editor_context: EditorContext = EditorContext{},
 stderr_store: ?*Blob.Store = null,
 stdin_store: ?*Blob.Store = null,
 stdout_store: ?*Blob.Store = null,
-websocket_mask: WebSocketClientMask = WebSocketClientMask{},
 
-uuid_entropy_cache: ?*EntropyCache = null,
+entropy_cache: ?*EntropyCache = null,
 
 // TODO: make this per JSGlobalObject instead of global
 // This does not handle ShadowRealm correctly!
@@ -25,21 +24,21 @@ tail_cleanup_hook: ?*CleanupHook = null,
 cleanup_hook: ?*CleanupHook = null,
 
 pub fn nextUUID(this: *RareData) [16]u8 {
-    if (this.uuid_entropy_cache == null) {
-        this.uuid_entropy_cache = default_allocator.create(EntropyCache) catch unreachable;
-        this.uuid_entropy_cache.?.init();
+    if (this.entropy_cache == null) {
+        this.entropy_cache = default_allocator.create(EntropyCache) catch unreachable;
+        this.entropy_cache.?.init();
     }
 
-    return this.uuid_entropy_cache.?.get();
+    return this.entropy_cache.?.get();
 }
 
 pub fn entropySlice(this: *RareData, len: usize) []u8 {
-    if (this.uuid_entropy_cache == null) {
-        this.uuid_entropy_cache = default_allocator.create(EntropyCache) catch unreachable;
-        this.uuid_entropy_cache.?.init();
+    if (this.entropy_cache == null) {
+        this.entropy_cache = default_allocator.create(EntropyCache) catch unreachable;
+        this.entropy_cache.?.init();
     }
 
-    return this.uuid_entropy_cache.?.slice(len);
+    return this.entropy_cache.?.slice(len);
 }
 
 pub const EntropyCache = struct {
diff --git a/src/http/websocket_http_client.zig b/src/http/websocket_http_client.zig
index 53caf36d9..6c9cc6004 100644
--- a/src/http/websocket_http_client.zig
+++ b/src/http/websocket_http_client.zig
@@ -477,46 +477,9 @@ pub fn NewHTTPUpgradeClient(comptime ssl: bool) type {
 }
 
 pub const Mask = struct {
-    const rand_buffer_size = 64;
-    bytes: [rand_buffer_size]u8 = [_]u8{0} ** rand_buffer_size,
-    offset: usize = 0,
-    needs_reset: bool = true,
-    queued_reset: bool = false,
-
-    pub fn from(globalObject: *JSC.JSGlobalObject) *Mask {
-        return &globalObject.bunVM().rareData().websocket_mask;
-    }
-
-    pub fn get(this: *Mask) [4]u8 {
-        if (this.needs_reset) {
-            this.needs_reset = false;
-            this.offset = 0;
-            std.crypto.random.bytes(&this.bytes);
-        }
-        const offset = this.offset;
-        const wrapped = offset % rand_buffer_size;
-        var mask = this.bytes[wrapped..][0..4].*;
-        if (offset > rand_buffer_size) {
-            const wrapped2 = @truncate(u8, wrapped);
-            mask[0] +%= wrapped2;
-            mask[1] +%= wrapped2;
-            mask[2] +%= wrapped2;
-            mask[3] +%= wrapped2;
-        }
-
-        this.offset += 4;
-
-        if (!this.queued_reset and this.offset % rand_buffer_size == 0) {
-            this.queued_reset = true;
-            uws.Loop.get().?.nextTick(*Mask, this, reset);
-        }
-
-        return mask;
-    }
-
-    pub fn fill(this: *Mask, mask_buf: *[4]u8, output_: []u8, input_: []const u8) void {
-        const mask = this.get();
-        mask_buf.* = mask;
+    pub fn fill(globalThis: *JSC.JSGlobalObject, mask_buf: *[4]u8, output_: []u8, input_: []const u8) void {
+        mask_buf.* = globalThis.bunVM().rareData().entropySlice(4)[0..4].*;
+        const mask = mask_buf.*;
 
         const skip_mask = @bitCast(u32, mask) == 0;
         if (!skip_mask) {
@@ -594,11 +557,6 @@ pub const Mask = struct {
             }
         }
     }
-
-    pub fn reset(this: *Mask) void {
-        this.queued_reset = false;
-        this.needs_reset = true;
-    }
 };
 
 const ReceiveState = enum {
@@ -750,7 +708,7 @@ const Copy = union(enum) {
                 header.opcode = Opcode.Text;
                 header.writeHeader(std.io.fixedBufferStream(buf).writer(), encode_into_result.written) catch unreachable;
 
-                Mask.from(globalThis).fill(buf[mask_offset..][0..4], to_mask[0..content_byte_len], to_mask[0..content_byte_len]);
+                Mask.fill(globalThis, buf[mask_offset..][0..4], to_mask[0..content_byte_len], to_mask[0..content_byte_len]);
             },
             .latin1 => |latin1| {
                 const encode_into_result = strings.copyLatin1IntoUTF8(to_mask, []const u8, latin1);
@@ -759,13 +717,13 @@ const Copy = union(enum) {
                 header.len = WebsocketHeader.packLength(encode_into_result.written);
                 header.opcode = Opcode.Text;
                 header.writeHeader(std.io.fixedBufferStream(buf).writer(), encode_into_result.written) catch unreachable;
-                Mask.from(globalThis).fill(buf[mask_offset..][0..4], to_mask[0..content_byte_len], to_mask[0..content_byte_len]);
+                Mask.fill(globalThis, buf[mask_offset..][0..4], to_mask[0..content_byte_len], to_mask[0..content_byte_len]);
             },
             .bytes => |bytes| {
                 header.len = WebsocketHeader.packLength(bytes.len);
                 header.opcode = Opcode.Binary;
                 header.writeHeader(std.io.fixedBufferStream(buf).writer(), bytes.len) catch unreachable;
-                Mask.from(globalThis).fill(buf[mask_offset..][0..4], to_mask[0..content_byte_len], bytes);
+                Mask.fill(globalThis, buf[mask_offset..][0..4], to_mask[0..content_byte_len], bytes);
             },
             .raw => unreachable,
         }
@@ -1259,7 +1217,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
             this.ping_frame_bytes[0..2].* = @bitCast(u16, header);
 
             if (to_mask.len > 0) {
-                Mask.from(this.globalThis).fill(this.ping_frame_bytes[2..6], to_mask, to_mask);
+                Mask.fill(this.globalThis, this.ping_frame_bytes[2..6], to_mask, to_mask);
                 return this.enqueueEncodedBytes(socket, this.ping_frame_bytes[0 .. 6 + @as(usize, this.ping_len)]);
             } else {
                 return this.enqueueEncodedBytes(socket, this.ping_frame_bytes[0..2]);
@@ -1296,7 +1254,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
 
             // we must mask the code
             var slice = final_body_bytes[0..(8 + body_len)];
-            Mask.from(this.globalThis).fill(mask_buf, slice[6..], slice[6..]);
+            Mask.fill(this.globalThis, mask_buf, slice[6..], slice[6..]);
 
             if (this.enqueueEncodedBytesMaybeFinal(socket, slice, true)) {
                 this.dispatchClose();