Building a UDP-based API

In my last post I introduced the concept of treating UDP network packets as a serverless event source. In this post I explore building a UDP-based application programming interfaces (APIs) with a request-response model.

Background

The modern web is brim-full of HTTP APIs, and we have excellent tooling to simplify and de-risk building them.

The simplicity means we don't need to write network stacks, web servers or parsers: we get them as part of the library or platform. Writing an HTTP server is as simple as:

from http.server import HTTPServer, BaseHTTPRequestHandler

class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.end_headers()
        self.wfile.write(b"OK")

HTTPServer(("", 8080), Handler).serve_forever()

If we want to get even simpler, we can use API Gateway and write a Lambda function to avoid paying for an idle container or instance to run our server. In that case the handler may be something like:

def handler(event, context):
    return {
        "statusCode": 200,
        "body": "OK"
    }

With this approach we have no costs for reserved capacity and no infrastructure to manage. This is what I mean by de-risking: We can build such APIs all day long and pay only when (if) they're actually used. Use rate limiting and WAF for abuse protection and reduce the risk even further.

What about UDP then?

A UDP API

A server for a UDP API is very similar. We're going to write it to listen for incoming packets on any port to keep with the example from the previous article.

In this case we want to support responses as an API should. We can't use EventBridge since it has a write-only API (responses relate to the success of the write; nothing custom).

So we're going to use Lambda instead and send what it returns as the response:

import socket
import struct
import boto3
import json
import base64

raw_sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_UDP)
lambda_client = boto3.client("lambda")
reply_socks = {}

while True:
    data, addr = raw_sock.recvfrom(65535)

    ip_header_len = (data[0] & 0x0F) * 4
    src_port, dst_port = struct.unpack('!HH', data[ip_header_len:ip_header_len + 4])
    payload = data[ip_header_len + 8:]
    src_ip = addr[0]

    response = lambda_client.invoke(
        FunctionName="my-function",
        Payload=json.dumps({
            "Port": dst_port,
            "Data": base64.b64encode(payload).decode()
        }),
    )
    reply = json.loads(response["Payload"].read())

    if reply:
        if dst_port not in reply_socks:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.bind(("", dst_port))
            reply_socks[dst_port] = s

        reply_socks[dst_port].sendto(reply.encode(), (src_ip, src_port))

Note that we're keeping a map of reply sockets to ensure that our replies appear to be "from" the same port the incoming packet has as the destination (the bind() call locks the socket to sending from a single port). If we didn't do this our replies would probably be ignored (or blocked by a firewall) since the client expects to have a two-way conversation with the given server port.

Now we can implement my-function however we'd like, including port-specific handling:

import json

def handler(event, context):
    if event.get("Port") == 4321:
        return "OK"

But we really want a more "API Gateway like" experience, and we can get there. The Constrained Application Protocol (CoAP) is the UDP-based cousin of HTTP. It has paths, methods, etc. and the basis for good general purpose UDP APIs.

So let's add support for CoAP to the server:

import socket
import struct
import boto3
import json
import base64
from aiocoap import Message, ACK, CONTENT

raw_sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_UDP)
lambda_client = boto3.client("lambda")
reply_socks = {}

ROUTES = {
    4321: {"function": "my-function",   "protocol": "raw"},
    5678: {"function": "coap-function", "protocol": "coap"},
}

def parse_coap(data):
    msg = Message.decode(data)
    return {
        "code": str(msg.code),
        "message_id": msg.mid,
        "token": base64.b64encode(msg.token).decode(),
        "path": "/" + "/".join(msg.opt.uri_path),
        "payload": base64.b64encode(msg.payload).decode(),
    }

def build_coap_ack(coap, body):
    return Message(
        mtype=ACK,
        code=CONTENT,
        mid=coap["message_id"],
        token=base64.b64decode(coap["token"]),
        payload=body.encode()
    ).encode()

while True:
    data, addr = raw_sock.recvfrom(65535)
    ip_header_len = (data[0] & 0x0F) * 4
    src_port, dst_port = struct.unpack('!HH', data[ip_header_len:ip_header_len + 4])
    payload = data[ip_header_len + 8:]
    src_ip = addr[0]

    route = ROUTES.get(dst_port)
    if not route:
        continue

    if route["protocol"] == "coap":
        coap = parse_coap(payload)
        lambda_payload = json.dumps(coap)
    else:
        lambda_payload = json.dumps({"Port": dst_port, "Data": base64.b64encode(payload).decode()})

    response = lambda_client.invoke(FunctionName=route["function"], Payload=lambda_payload)
    reply = json.loads(response["Payload"].read())

    if reply:
        if dst_port not in reply_socks:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.bind(("", dst_port))
            reply_socks[dst_port] = s

        packet = build_coap_ack(coap, reply) if route["protocol"] == "coap" else reply.encode()
        reply_socks[dst_port].sendto(packet, (src_ip, src_port))

In reality the configuration would be external to the server, but we're keeping things simple here. To add a new CoAP API endpoint we can just add to the configuration. To implement a CoAP handler, we build a simple Lambda that just processes the JSON object payload.

For example, here is a Lambda implementing a "time" service:

import datetime

def handler(event, context):
    if event.get("path") == "/time":
        return datetime.datetime.now().isoformat()

Adding more paths and methods to the Lambda will feel very much like building an HTTP handler.

Summary

This is the start of an interesting architecture that decouples network listening and replies from handling the packet data:

• We can add new paths and methods to the CoAP Lambda without touching the server at all.
• We can add new ports and custom protocols to the server in a clean, configuration driven way.
• Using CoAP, we're within sight of the simplicity of the tooling for HTTP but lack the ease of configuration and features.
• We're running a server, and hence haven't really de-risked as much as we'd like.

To that last point, we can't quite match API Gateway approach (there is no AWS service that supports UDP in this way, but Proxylity does).