gitweb.ps.run Git - chirp/blob - src/main.zig

   1 const std = @import("std");
   2 const lmdb = @import("lmdb");
   3 const http = @import("http");
   4
   5 // db {{{
   6
   7 const Prng = struct {
   8     var prng: std.Random.DefaultPrng = std.Random.DefaultPrng.init(0);
   9
  10     pub fn gen_id(dbi: anytype) Id {
  11         var id = Prng.prng.next();
  12
  13         while (dbi.has(id)) {
  14             id = Prng.prng.next();
  15         }
  16
  17         return id;
  18     }
  19
  20     pub fn gen_str(dbi: anytype, comptime len: usize) [len]u8 {
  21         var buf: [len / 2]u8 = undefined;
  22         var res: [len]u8 = undefined;
  23         Prng.prng.fill(&buf);
  24         for (0..len / 2) |i| {
  25             res[i * 2 + 0] = 'a' + (buf[i] % 16);
  26             res[i * 2 + 1] = 'a' + (buf[i] >> 4 % 16);
  27         }
  28
  29         while (dbi.has(res)) {
  30             Prng.prng.fill(&buf);
  31             for (0..len / 2) |i| {
  32                 res[i * 2 + 0] = 'a' + (buf[i] % 16);
  33                 res[i * 2 + 1] = 'a' + (buf[i] >> 4 % 16);
  34             }
  35         }
  36
  37         return res;
  38     }
  39 };
  40
  41 const Db = struct {
  42     fn users(txn: *const lmdb.Txn) !lmdb.Dbi(Id, User) {
  43         return try txn.dbi("users", Id, User);
  44     }
  45     fn user_ids(txn: *const lmdb.Txn) !lmdb.Dbi(Username, Id) {
  46         return try txn.dbi("user_ids", Username, Id);
  47     }
  48     fn sessions(txn: *const lmdb.Txn) !lmdb.Dbi(SessionToken, Id) {
  49         return try txn.dbi("sessions", SessionToken, Id);
  50     }
  51 };
  52
  53 // }}}
  54
  55 // content {{{
  56
  57 const User = struct {
  58     // TODO: choose sizes
  59     username: Username,
  60     password_hash: PasswordHash,
  61 };
  62
  63 const SessionTokenLen = 16;
  64
  65 const Id = u64;
  66 const Username = std.BoundedArray(u8, 16);
  67 const PasswordHash = std.BoundedArray(u8, 128);
  68 const SessionToken = [SessionTokenLen]u8;
  69 const CookieValue = std.BoundedArray(u8, 128);
  70
  71 pub fn hash_password(password: []const u8) !PasswordHash {
  72     var hash_buffer = try PasswordHash.init(128);
  73
  74     // TODO: choose buffer size
  75     // TODO: dont allocate on stack, maybe zero memory?
  76     var buffer: [1024 * 10]u8 = undefined;
  77     var alloc = std.heap.FixedBufferAllocator.init(&buffer);
  78
  79     // TODO: choose limits
  80     const result = try std.crypto.pwhash.argon2.strHash(password, .{
  81         .allocator = alloc.allocator(),
  82         .params = std.crypto.pwhash.argon2.Params.fromLimits(1000, 1024),
  83     }, hash_buffer.slice());
  84
  85     try hash_buffer.resize(result.len);
  86
  87     return hash_buffer;
  88 }
  89
  90 pub fn verify_password(password: []const u8, hash: PasswordHash) bool {
  91     var buffer: [1024 * 10]u8 = undefined;
  92     var alloc = std.heap.FixedBufferAllocator.init(&buffer);
  93
  94     if (std.crypto.pwhash.argon2.strVerify(hash.constSlice(), password, .{
  95         .allocator = alloc.allocator(),
  96     })) {
  97         return true;
  98     } else |err| {
  99         std.debug.print("verify error: {}\n", .{err});
 100         return false;
 101     }
 102 }
 103
 104 pub fn register_user(env: *lmdb.Env, username: []const u8, password: []const u8) !bool {
 105     const username_array = try Username.fromSlice(username);
 106
 107     const txn = try env.txn();
 108     defer {
 109         txn.commit();
 110         env.sync();
 111     }
 112
 113     const users = try Db.users(&txn);
 114     const user_ids = try Db.user_ids(&txn);
 115
 116     if (user_ids.has(username_array)) {
 117         return false;
 118     } else {
 119         const user_id = Prng.gen_id(users);
 120         users.put(user_id, User{
 121             .username = username_array,
 122             .password_hash = try hash_password(password),
 123         });
 124
 125         user_ids.put(username_array, user_id);
 126
 127         return true;
 128     }
 129 }
 130
 131 pub fn login_user(env: *lmdb.Env, username: []const u8, password: []const u8) !SessionToken {
 132     const username_array = try Username.fromSlice(username);
 133
 134     const txn = try env.txn();
 135     defer {
 136         txn.commit();
 137         env.sync();
 138     }
 139
 140     const user_ids = try Db.user_ids(&txn);
 141     const user_id = user_ids.get(username_array) orelse return error.UnknownUsername;
 142     std.debug.print("id: {}\n", .{user_id});
 143
 144     const users = try Db.users(&txn);
 145     if (users.get(user_id)) |user| {
 146         if (verify_password(password, user.password_hash)) {
 147             const sessions = try Db.sessions(&txn);
 148             const session_token = Prng.gen_str(sessions, SessionTokenLen);
 149             sessions.put(session_token, user_id);
 150             return session_token;
 151         } else {
 152             return error.IncorrectPassword;
 153         }
 154     } else {
 155         return error.UserNotFound;
 156     }
 157 }
 158
 159 fn logout_user(env: *lmdb.Env, session_token: SessionToken) !void {
 160     const txn = try env.txn();
 161     defer {
 162         txn.commit();
 163         env.sync();
 164     }
 165
 166     const sessions = try Db.sessions(&txn);
 167     sessions.del(session_token);
 168 }
 169
 170 fn get_session_user(env: *lmdb.Env, session_token: SessionToken) !User {
 171     const txn = try env.txn();
 172     defer txn.abort();
 173
 174     const sessions = try Db.sessions(&txn);
 175     const users = try Db.users(&txn);
 176
 177     if (sessions.get(session_token)) |user_id| {
 178         return users.get(user_id) orelse error.UnknownUser;
 179     } else {
 180         return error.SessionNotFound;
 181     }
 182 }
 183
 184 // }}}
 185
 186 fn list_users(env: *lmdb.Env) !void {
 187     const txn = try env.txn();
 188     defer txn.abort();
 189
 190     // const users = try Db.users(&txn);
 191     const users = try txn.dbi("users", Id, User);
 192     var cursor = try users.cursor();
 193
 194     var key: Id = undefined;
 195     var user_maybe = cursor.get(&key, .First);
 196
 197     while (user_maybe) |*user| {
 198         std.debug.print("[{}] {s}\n", .{ key, user.username.constSlice() });
 199
 200         user_maybe = cursor.get(&key, .Next);
 201     }
 202 }
 203 fn list_user_ids(env: *lmdb.Env) !void {
 204     const txn = try env.txn();
 205     defer txn.abort();
 206
 207     const user_ids = try Db.user_ids(&txn);
 208     var cursor = try user_ids.cursor();
 209
 210     var key: Username = undefined;
 211     var user_id_maybe = cursor.get(&key, .First);
 212
 213     while (user_id_maybe) |user_id| {
 214         std.debug.print("[{s}] {}\n", .{ key.constSlice(), user_id });
 215
 216         user_id_maybe = cursor.get(&key, .Next);
 217     }
 218 }
 219
 220 fn list_sessions(env: *lmdb.Env) !void {
 221     const txn = try env.txn();
 222     defer txn.abort();
 223
 224     const sessions = try Db.sessions(&txn);
 225     var cursor = try sessions.cursor();
 226
 227     var key: SessionToken = undefined;
 228     var user_id_maybe = cursor.get(&key, .First);
 229
 230     while (user_id_maybe) |user_id| {
 231         std.debug.print("[{s}] {}\n", .{ key, user_id });
 232
 233         user_id_maybe = cursor.get(&key, .Next);
 234     }
 235 }
 236
 237 const ReqBufferSize = 4096;
 238 const ResHeadBufferSize = 4096;
 239 const ResBodyBufferSize = 4096;
 240
 241 pub fn main() !void {
 242     // server
 243     var server = try http.Server.init("::", 8080);
 244     defer server.deinit();
 245
 246     // lmdb
 247     var env = lmdb.Env.open("db", 1024 * 100);
 248     defer env.close();
 249
 250     std.debug.print("Users:\n", .{});
 251     try list_users(&env);
 252     std.debug.print("User IDs:\n", .{});
 253     try list_user_ids(&env);
 254     std.debug.print("Sessions:\n", .{});
 255     try list_sessions(&env);
 256
 257     try handle_connection(&server, &env);
 258     // const ThreadCount = 1;
 259     // var ts: [ThreadCount]std.Thread = undefined;
 260
 261     // for (0..ThreadCount) |i| {
 262     //     ts[i] = try std.Thread.spawn(.{}, handle_connection, .{ &server, &env });
 263     // }
 264     // for (0..ThreadCount) |i| {
 265     //     ts[i].join();
 266     // }
 267
 268     std.debug.print("done\n", .{});
 269 }
 270
 271 fn handle_connection(server: *http.Server, env: *lmdb.Env) !void {
 272     // TODO: static?
 273     var req_buffer: [ReqBufferSize]u8 = undefined;
 274     var res_head_buffer: [ResHeadBufferSize]u8 = undefined;
 275     var res_body_buffer: [ResBodyBufferSize]u8 = undefined;
 276
 277     accept: while (true) {
 278         server.wait();
 279
 280         while (try server.next_request(&req_buffer)) |req| {
 281             // std.debug.print("[{}]: {s}\n", .{ req.method, req.target });
 282
 283             // reponse
 284             var res = http.Response.init(req.fd, &res_head_buffer, &res_body_buffer);
 285
 286             // check session token
 287             var logged_in: ?struct {
 288                 user: User,
 289                 session_token: SessionToken,
 290             } = null;
 291
 292             if (req.get_cookie("session_token")) |session_token_str| {
 293                 var session_token: SessionToken = undefined;
 294                 std.mem.copyForwards(u8, &session_token, session_token_str);
 295                 // const session_token = try std.fmt.parseUnsigned(SessionToken, session_token_str, 10);
 296                 // const session_token = std.mem.bytesToValue(SessionToken, session_token_str);
 297                 if (get_session_user(env, session_token)) |user| {
 298                     logged_in = .{
 299                         .user = user,
 300                         .session_token = session_token,
 301                     };
 302                 } else |err| {
 303                     std.debug.print("get_session_user err: {}\n", .{err});
 304
 305                     try res.add_header(
 306                         "Set-Cookie",
 307                         .{"session_token=deleted; Expires=Thu, 01 Jan 1970 00:00:00 GMT"},
 308                     );
 309                 }
 310             }
 311
 312             // html
 313             if (req.method == .GET) {
 314                 if (std.mem.eql(u8, req.target, "/register")) {
 315                     try res.write(
 316                         \\<form action="/register" method="post">
 317                         \\<input type="text" name="username" />
 318                         \\<input type="password" name="password" />
 319                         \\<input type="submit" value="Register" />
 320                         \\</form>
 321                     , .{});
 322                     try res.send();
 323                 } else if (std.mem.eql(u8, req.target, "/login")) {
 324                     try res.write(
 325                         \\<form action="/login" method="post">
 326                         \\<input type="text" name="username" />
 327                         \\<input type="password" name="password" />
 328                         \\<input type="submit" value="Login" />
 329                         \\</form>
 330                     , .{});
 331                     try res.send();
 332                 } else {
 333                     if (logged_in) |login| {
 334                         try res.write(
 335                             \\<a href="/user/{s}">Home</a>
 336                             \\<form action="/logout" method="post"><input type="submit" value="Logout" /></form>
 337                             \\<form action="/quit" method="post"><input type="submit" value="Quit" /></form>
 338                         , .{login.user.username.constSlice()});
 339                         try res.send();
 340                     } else {
 341                         try res.write(
 342                             \\<a href="/register">Register</a>
 343                             \\<a href="/login">Login</a>
 344                             \\<form action="/quit" method="post"><input type="submit" value="Quit" /></form>
 345                         , .{});
 346                         try res.send();
 347                     }
 348                 }
 349             }
 350             // api
 351             else {
 352                 if (std.mem.eql(u8, req.target, "/register")) {
 353                     // TODO: handle args not supplied
 354                     const username = req.get_value("username").?;
 355                     const password = req.get_value("password").?;
 356
 357                     std.debug.print("New user: {s} {s}\n", .{ username, password });
 358                     if (try register_user(env, username, password)) {
 359                         try res.redirect("/login");
 360                     } else {
 361                         try res.redirect("/register");
 362                     }
 363                     try res.send();
 364                 } else if (std.mem.eql(u8, req.target, "/login")) {
 365                     // TODO: handle args not supplied
 366                     const username = req.get_value("username").?;
 367                     const password = req.get_value("password").?;
 368
 369                     std.debug.print("New login: {s} {s}\n", .{ username, password });
 370                     if (login_user(env, username, password)) |session_token| {
 371                         res.status = .see_other;
 372                         try res.add_header(
 373                             "Location",
 374                             .{ "/user/{s}", .{username} },
 375                         );
 376                         try res.add_header(
 377                             "Set-Cookie",
 378                             .{ "session_token={s}; Secure; HttpOnly", .{session_token} },
 379                         );
 380
 381                         try res.send();
 382                     } else |err| {
 383                         std.debug.print("login_user err: {}\n", .{err});
 384                         try res.redirect("/login");
 385                         try res.send();
 386                     }
 387                 } else if (std.mem.eql(u8, req.target, "/logout")) {
 388                     if (logged_in) |login| {
 389                         try logout_user(env, login.session_token);
 390
 391                         try res.add_header(
 392                             "Set-Cookie",
 393                             .{"session_token=deleted; Expires=Thu, 01 Jan 1970 00:00:00 GMT"},
 394                         );
 395
 396                         try res.redirect("/");
 397                         try res.send();
 398                     }
 399                 } else if (std.mem.eql(u8, req.target, "/quit")) {
 400                     try res.redirect("/");
 401                     try res.send();
 402                     break :accept;
 403                 } else {
 404                     // try req.respond(
 405                     //     \\<p>POST</p>
 406                     // , .{});
 407                     try res.write("<p>{s}</p>", .{req.target});
 408                     try res.send();
 409                 }
 410             }
 411         }
 412     }
 413 }