Final Project: A Link Shortener

[package]
name = "links"
description = "A small link shortener"

[grants]
listen = ["127.0.0.1"]
net = ["127.0.0.1"]

// Short codes for links, and checks on the URLs people send. Everything
// here is pure: no capabilities, so it can only compute.

fn alphabet() -> List[Str] {
    "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ".chars()
}

// The code for the link with this id: the id written in base 62.
pub fn encode(id: Int) -> Str {
    let digits = alphabet()
    let mut n = id
    let mut code = ""
    while true {
        code = digits[n % 62] + code
        n = n / 62
        if n == 0 {
            break
        }
    }
    code
}

// The id a code stands for, or None if it is not a code.
pub fn decode(code: Str) -> Option[Int] {
    if code.is_empty() || code.len() > 10 {
        return None
    }
    let digits = alphabet()
    let mut n = 0
    for c in code.chars() {
        let d = digits.find_index(c)?
        n = n * 62 + d
    }
    Some(n)
}

fn find_index(self: List[Str], x: Str) -> Option[Int] {
    for (i, y) in self.with_index() {
        if y == x {
            return Some(i)
        }
    }
    None
}

// Whether a URL is one we are willing to redirect to.
pub fn valid_url(url: Str) -> Bool {
    let rest = if url.starts_with("https://") {
        url.slice(8, url.len())
    } else if url.starts_with("http://") {
        url.slice(7, url.len())
    } else {
        return false
    }
    let host = rest.split("/")[0]
    host.contains(".") && !host.contains(" ") && url.len() <= 2000
}

test "encode" for (id, want) in [(0, "0"), (61, "Z"), (62, "10"), (125, "21")] {
    assert_eq(encode(id), want)
}

test "decode undoes encode" {
    for id in [0, 1, 61, 62, 3843, 1_000_000] {
        assert_eq(decode(encode(id)), Some(id))
    }
}

test "decode rejects what is not a code" {
    assert_eq(decode(""), None)
    assert_eq(decode("ab-c"), None)
}

test "valid_url" for (url, ok) in [
    ("https://example.com/a?b=c", true),
    ("http://example.com", true),
    ("ftp://example.com", false),
    ("https://localhost", false),
    ("javascript:alert(1)", false),
] {
    assert_eq(valid_url(url), ok)
}

// Where links live: a SQLite database. Every function takes the Db and
// nothing else, so this module can touch that database and no other thing.
use sql
use codes

pub struct Link {
    code: Str,
    url: Str,
    hits: Int,
}

pub fn setup(db: sql.Db) -> Result[(), IoError] {
    db.exec("CREATE TABLE IF NOT EXISTS links (id INTEGER PRIMARY KEY, url TEXT NOT NULL, hits INTEGER NOT NULL DEFAULT 0)", [])?
    Ok(())
}

pub fn add(db: sql.Db, url: Str) -> Result[Link, IoError] {
    let row = db.query_one("INSERT INTO links (url) VALUES (?) RETURNING id", [sql.Text(url)])?
    let id = row?.int("id") ?? 0
    Ok(Link { code: codes.encode(id), url, hits: 0 })
}

pub fn find(db: sql.Db, code: Str) -> Result[Option[Link], IoError] {
    let id = match codes.decode(code) {
        Some(id) => id,
        None => {
            return Ok(None)
        },
    }
    let row = db.query_one("SELECT url, hits FROM links WHERE id = ?", [sql.Int(id)])?
    Ok(row.map(|r| Link { code, url: r.text("url") ?? "", hits: r.int("hits") ?? 0 }))
}

// Follows a link: counts the visit and returns where it goes.
pub fn visit(db: sql.Db, code: Str) -> Result[Option[Str], IoError] {
    let id = match codes.decode(code) {
        Some(id) => id,
        None => {
            return Ok(None)
        },
    }
    let row = db.query_one("UPDATE links SET hits = hits + 1 WHERE id = ? RETURNING url", [sql.Int(id)])?
    Ok(row?.text("url"))
}

// The HTTP side: requests in, responses out. The handler gets the database
// and nothing else, whatever a request asks for.
use http
use json
use sql
use codes
use store

pub fn handle(db: sql.Db, r: http.Request) -> http.Response {
    let path = r.path.split("?")[0]
    let parts = path.split("/").filter(|p| p != "")
    let result = match (r.method, parts) {
        ("POST", ["links"]) => create(db, r.body),
        ("GET", ["links", code]) => stats(db, code),
        ("GET", [code]) => follow(db, code),
        ("GET", []) => Ok(http.text(200, "POST /links with \{\"url\": \"...\"\} to shorten a link\n")),
        _ => Ok(http.text(404, "not found\n")),
    }
    match result {
        Ok(resp) => resp,
        Err(e) => http.text(500, "storage error: {e}\n"),
    }
}

fn create(db: sql.Db, body: Str) -> Result[http.Response, IoError] {
    let url = match json.parse(body) {
        Ok(v) => v.get("url")?.as_str() ?? "",
        Err(_) => "",
    }
    if !codes.valid_url(url) {
        return Ok(http.text(400, "send \{\"url\": \"https://...\"\}\n"))
    }
    let link = store.add(db, url)?
    Ok(http.json(201, link_json(link)))
}

fn stats(db: sql.Db, code: Str) -> Result[http.Response, IoError] {
    match store.find(db, code)? {
        Some(link) => Ok(http.json(200, link_json(link))),
        None => Ok(http.text(404, "no link {code}\n")),
    }
}

fn follow(db: sql.Db, code: Str) -> Result[http.Response, IoError] {
    match store.visit(db, code)? {
        Some(url) => Ok(http.Response { status: 302, headers: [http.Header { name: "Location", value: url }], body: "" }),
        None => Ok(http.text(404, "no link {code}\n")),
    }
}

fn link_json(l: store.Link) -> Str {
    json.stringify(json.Obj([
        json.field("code", json.Str(l.code)),
        json.field("url", json.Str(l.url)),
        json.field("hits", json.Num(l.hits.to_float())),
    ]))
}

// ---------- tests: the handler is a plain function, so tests call it ----------

fn request(db: sql.Db, method: Str, path: Str, body: Str) -> http.Response {
    handle(db, http.Request { method, path, version: "HTTP/1.1", headers: [], body })
}

fn test_db() -> sql.Db {
    let db = sql.memory() ?? panic("no database")
    store.setup(db) ?? panic("no table")
    db
}

test "shorten, follow, count" {
    let db = test_db()
    let made = request(db, "POST", "/links", "\{\"url\": \"https://example.com/a\"\}")
    assert_eq(made.status, 201)
    assert_eq(made.body, "\{\"code\":\"1\",\"url\":\"https://example.com/a\",\"hits\":0\}")

    let go = request(db, "GET", "/1", "")
    assert_eq(go.status, 302)
    assert_eq(go.header("location"), Some("https://example.com/a"))

    assert_eq(request(db, "GET", "/links/1", "").body.contains("\"hits\":1"), true)
}

test "refuses bad input" for (method, path, body, status) in [
    ("POST", "/links", "\{\"url\": \"javascript:alert(1)\"\}", 400),
    ("POST", "/links", "not json", 400),
    ("GET", "/zzzz", "", 404),
    ("GET", "/links/!", "", 404),
    ("DELETE", "/links", "", 404),
] {
    assert_eq(request(test_db(), method, path, body).status, status)
}

// A link shortener.
//
//   tessera run main.tsr                  a demo: server and client in one run
//   tessera run main.tsr -- serve 8080    serve on 127.0.0.1:8080
use http
use sql
use store
use api

fn main(root: Root) -> Result[(), dyn Error] {
    let db = sql.memory()? "opening the database"
    store.setup(db)? "creating the table"
    let args = root.args()
    if args.get(0) == Some("serve") {
        let port = args.get(1) ?? "8080"
        let l = root.listen.on("127.0.0.1:{port}")?
        root.stdout.print("serving on http://127.0.0.1:{port}")
        return Err(http.serve_forever(l, |r| api.handle(db, r)))
    }
    demo(root.listen, root.net.only("127.0.0.1"), root.stdout, db)
}

// Starts the server on a free port and talks to it.
fn demo(listen: Listen, net: Net, out: Stdout, db: sql.Db) -> Result[(), dyn Error] {
    let l = listen.on("127.0.0.1:0")?
    let base = "http://127.0.0.1:{l.port()}"
    let lines = group g {
        g.go(|| {
            let _ = http.serve(l, 4, |r| api.handle(db, r))
        })
        g.go(|| client(net, base))
    }
    for line in lines[0] {
        out.print(line)
    }
    Ok(())
}

fn client(net: Net, base: Str) -> List[Str] {
    let mut log: List[Str] = []
    let steps = [
        ("POST", "/links", "\{\"url\": \"https://example.com/a/long/page\"\}"),
        ("GET", "/1", ""),
        ("GET", "/links/1", ""),
        ("GET", "/nope", ""),
    ]
    for (method, path, body) in steps {
        let headers = [http.Header { name: "Content-Type", value: "application/json" }]
        match http.request(net, method, base + path, headers, body) {
            Ok(r) => {
                let to = r.header("location").map(|u| " -> {u}") ?? ""
                log.push("{method} {path}: {r.status}{to} {r.body}".trim())
            },
            Err(e) => log.push("{method} {path}: {e}"),
        }
    }
    log
}