fjord/server.odin

package fjord

import "base:runtime"
import "core:bufio"
import "core:bytes"
import "core:io"
import "core:log"
import "core:net"
import "core:path/slashpath"
import "core:strconv"
import "core:strings"
import "core:sync"
import "core:thread"
import "core:unicode"

import http "http"

Server :: struct($Error_Type: typeid) {
    endpoint:          net.Endpoint,
    not_found_handler: #type proc(
        request: ^http.Request,
    ) -> (
        http.Response,
        Error_Type,
    ),
    running:           bool,
    router:            Router(Error_Type),
    thread_pool:       thread.Pool,
    thread_count:      uint,
    allocator:         runtime.Allocator,
}

@(private)
read_connection :: proc(
    client_socket: net.TCP_Socket,
    allocator := context.temp_allocator,
) -> (
    data: []byte,
    err: http.RequestError,
) {
    TCP_CHUNK_SIZE :: 1024
    HTTP_END_SEQUENCE :: []byte{'\r', '\n', '\r', '\n'}

    buffer: bytes.Buffer
    bytes.buffer_init_allocator(&buffer, 0, 0, allocator)

    chunk: [TCP_CHUNK_SIZE]byte
    for {
        bytes_read, net_err := net.recv_tcp(client_socket, chunk[:])
        if net_err != nil do return data, .NetworkError

        if bytes_read == 0 do break

        prev_len := len(buffer.buf)
        bytes.buffer_write(&buffer, chunk[:bytes_read])

        search_start := max(0, prev_len - 3)
        if bytes.index(buffer.buf[search_start:], HTTP_END_SEQUENCE) != -1 do break
    }

    return buffer.buf[:], nil
}

server_init :: proc(
    server: ^Server($Error_Type),
    endpoint: net.Endpoint,
    thread_count: uint,
    allocator := context.allocator,
) {
    server^ = Server(Error_Type) {
        endpoint = endpoint,
        not_found_handler = proc(
            request: ^http.Request,
        ) -> (
            http.Response,
            Error_Type,
        ) {
            body := "<div>Not found</div>"
            http_response := http.make_response(
                .NotFound,
                transmute([]byte)body,
                .Html,
            )
            return http_response, nil
        },
        thread_count = thread_count,
        running = false,
        allocator = allocator,
    }

    thread.pool_init(&server.thread_pool, allocator, int(server.thread_count))
    thread.pool_start(&server.thread_pool)
    router_init(&server.router, allocator)
}

server_destroy :: proc(server: ^Server($Error_Type)) {
    if sync.atomic_load(&server.thread_pool.is_running) {
        thread.pool_join(&server.thread_pool)
    }

    thread.pool_destroy(&server.thread_pool)
    router_destroy(&server.router)
}

// This is used instead of default server_make param until compiler bug is fixed: https://github.com/odin-lang/Odin/issues/5792
server_set_not_found_handler :: proc(
    server: ^Server($Error_Type),
    not_found_handler: #type proc(
        request: ^http.Request,
    ) -> (
        http.Response,
        Error_Type,
    ),
) {
    server.not_found_handler = not_found_handler
}

server_add_route :: proc(
    server: ^Server($Error_Type),
    method: http.Method,
    path: []string,
    handler: #type proc(request: ^http.Request) -> (http.Response, Error_Type),
) {
    router_add_route(&server.router, method, path, handler)
}

server_remove_route :: proc(
    server: ^Server($Error_Type),
    method: http.Method,
    path: []string,
) -> (
    ok: bool,
) {
    return router_remove_route(&server.router, method, path)
}

server_shutdown :: proc(server: ^Server($Error_Type)) {
    sync.atomic_store(&server.running, false)
}

@(private)
serve :: proc(server: ^Server($Error_Type), client_socket: net.TCP_Socket) {
    data, read_err := read_connection(client_socket, context.temp_allocator)
    if read_err != nil {
        log.warnf("Failed to read request, reason: %s", read_err)
        return
    }

    http_request, unmarshall_err := http.unmarshall_request(
        data,
        context.temp_allocator,
    )
    if unmarshall_err != nil {
        log.warnf("Failed to unmarshall request, reason: %s", unmarshall_err)
        return
    }

    handler: #type proc(request: ^http.Request) -> (http.Response, Error_Type)
    ok: bool

    // todo: should sanitize better
    path := strings.split(http_request.path, "/", context.temp_allocator)[1:]

    handler, http_request.path_variables, ok = router_lookup(
        &server.router,
        http_request.method,
        path,
        context.temp_allocator,
    )
    if !ok do handler = server.not_found_handler

    http_response, handler_err := handler(&http_request)

    if handler_err != nil {
        log.warnf("Handler failed with error: %s", handler_err)
        return
    }

    marshalled_response := http.marshall_response(
        &http_response,
        context.temp_allocator,
    )
    net.send_tcp(client_socket, marshalled_response)
}

RequestThreadData :: struct($Error_Type: typeid) {
    server:        ^Server(Error_Type),
    client_socket: net.TCP_Socket,
}

@(private)
request_thread_task :: proc(t: thread.Task) {
    task_data := cast(^RequestThreadData(any))t.data
    defer free(task_data, t.allocator)

    serve(task_data.server, task_data.client_socket)

    net.close(task_data.client_socket)
    // context.temp_allocator is freed automatically on thread death.
}

@(private)
spawn_request_thread :: proc(
    server: ^Server($Error_Type),
    client_socket: net.TCP_Socket,
) {
    task_data := new(RequestThreadData(Error_Type))
    task_data^ = RequestThreadData(Error_Type) {
        server        = server,
        client_socket = client_socket,
    }

    // Not sure about sharing the default non temp allocator, have to think about
    thread.pool_add_task(
        &server.thread_pool,
        context.allocator,
        request_thread_task,
        rawptr(task_data),
    )
}

server_listen_and_serve :: proc(server: ^Server($Error_Type)) {
    server_socket, net_err := net.listen_tcp(server.endpoint)
    log.assert(net_err == nil, "Couldn't create TCP socket")
    defer net.close(server_socket)

    sync.atomic_store(&server.running, true)

    for sync.atomic_load(&server.running) {
        client_socket, _, net_err := net.accept_tcp(server_socket)
        if net_err != nil {
            log.warnf("Failed to accept TCP connection, reason: %s", net_err)
            continue
        }

        spawn_request_thread(server, client_socket)
    }

    thread.pool_finish(&server.thread_pool)
    thread.pool_join(&server.thread_pool)
}