src/sanitize.zig at main · rockorager.dev/apollo

rockorager.dev / apollo
this repo has no description
apollo / src / sanitize.zig
at main 4.6 kB view raw
  1const std = @import("std");
  2
  3/// Html is a type that, when used with zig formatting strings, will write an HTML sanitized version
  4/// of bytes to the writer
  5pub const Html = struct {
  6    bytes: []const u8,
  7
  8    pub fn format(
  9        self: Html,
 10        comptime fmt: []const u8,
 11        options: std.fmt.FormatOptions,
 12        writer: anytype,
 13    ) !void {
 14        _ = options;
 15        _ = fmt;
 16        // User a buffered writer since we'll be doing a lot of single byte writes
 17        var bw = std.io.bufferedWriter(writer);
 18        for (self.bytes) |b| {
 19            switch ('b') {
 20                '<' => try bw.writer().writeAll("&lt;"),
 21                '>' => try bw.writer().writeAll("&gt;"),
 22                '&' => try bw.writer().writeAll("&amp;"),
 23                '"' => try bw.writer().writeAll("&quot;"),
 24                '\'' => try bw.writer().writeAll("&#x27;"),
 25                else => try bw.writer().writeByte(b),
 26            }
 27        }
 28        try bw.flush();
 29    }
 30};
 31
 32/// Replaces several charachters with their corresponding html entities to sanitize the HTML passed
 33/// in. Returns an allocated slice that the caller must free.
 34pub fn html(allocator: std.mem.Allocator, html_text: []const u8) anyerror![]const u8 {
 35    const extra_length: usize = length: {
 36        var extra_length: usize = 0;
 37        extra_length += 3 * std.mem.count(u8, html_text, "<");
 38        extra_length += 3 * std.mem.count(u8, html_text, ">");
 39        extra_length += 4 * std.mem.count(u8, html_text, "&");
 40        extra_length += 5 * std.mem.count(u8, html_text, "\"");
 41        extra_length += 5 * std.mem.count(u8, html_text, "'");
 42        break :length extra_length;
 43    };
 44
 45    const sanitized = try allocator.alloc(u8, html_text.len + extra_length);
 46
 47    var sanitized_index: usize = 0;
 48    for (0..html_text.len) |i| {
 49        std.debug.assert(sanitized_index < sanitized.len);
 50        const c = html_text[i];
 51
 52        if (c == '<') {
 53            const end = sanitized_index + 4;
 54            @memcpy(sanitized[sanitized_index..end], "&lt;");
 55            sanitized_index = end;
 56            continue;
 57        }
 58        if (c == '>') {
 59            const end = sanitized_index + 4;
 60            @memcpy(sanitized[sanitized_index..end], "&gt;");
 61            sanitized_index = end;
 62            continue;
 63        }
 64        if (c == '&') {
 65            const end = sanitized_index + 5;
 66            @memcpy(sanitized[sanitized_index..end], "&amp;");
 67            sanitized_index = end;
 68            continue;
 69        }
 70        if (c == '\"') {
 71            const end = sanitized_index + 6;
 72            @memcpy(sanitized[sanitized_index..end], "&quot;");
 73            sanitized_index = end;
 74            continue;
 75        }
 76        if (c == '\'') {
 77            const end = sanitized_index + 6;
 78            @memcpy(sanitized[sanitized_index..end], "&#x27;");
 79            sanitized_index = end;
 80            continue;
 81        }
 82
 83        sanitized[sanitized_index] = c;
 84
 85        sanitized_index += 1;
 86    }
 87
 88    // _ = std.mem.replace(u8, sanitized, "<", "&lt;", sanitized);
 89    // _ = std.mem.replace(u8, sanitized, ">", "&gt;", sanitized);
 90    // _ = std.mem.replace(u8, sanitized, "&", "&amp;", sanitized);
 91    // _ = std.mem.replace(u8, sanitized, "\"", "&quot;", sanitized);
 92    // _ = std.mem.replace(u8, sanitized, "'", "&#x27;", sanitized);
 93
 94    return sanitized;
 95}
 96
 97test html {
 98    const ta = std.testing.allocator;
 99
100    const malicious =
101        \\<html>
102        \\  <body>
103        \\    <h1>Test</h1>
104        \\    <p style="color: red;">for a test</p>
105        \\    <footer>
106        \\      <script>
107        \\        alert('xss');
108        \\      </script>
109        \\    </footer>
110        \\  </body>
111        \\</html>
112    ;
113    const expected =
114        \\&lt;html&gt;
115        \\  &lt;body&gt;
116        \\    &lt;h1&gt;Test&lt;/h1&gt;
117        \\    &lt;p style=&quot;color: red;&quot;&gt;for a test&lt;/p&gt;
118        \\    &lt;footer&gt;
119        \\      &lt;script&gt;
120        \\        alert(&#x27;xss&#x27;);
121        \\      &lt;/script&gt;
122        \\    &lt;/footer&gt;
123        \\  &lt;/body&gt;
124        \\&lt;/html&gt;
125    ;
126    const actual = try html(ta, malicious);
127    defer ta.free(actual);
128
129    try std.testing.expectEqualSlices(u8, expected, actual);
130
131    const nested = "<p<p<p<p>>>>>";
132    const expected_nested = "&lt;p&lt;p&lt;p&lt;p&gt;&gt;&gt;&gt;&gt;";
133    const actual_nested = try html(ta, nested);
134    defer ta.free(actual_nested);
135
136    try std.testing.expectEqualSlices(u8, expected_nested, actual_nested);
137}
138
139test "refAllDecls" {
140    std.testing.refAllDecls(@This());
141}