diff --git a/lua/framebuffer.tl b/lua/framebuffer.tl
index 18edb00..94fc03d 100644
--- a/lua/framebuffer.tl
+++ b/lua/framebuffer.tl
@@ -136,4 +136,4 @@ local function wrap(parent: term.Redirect): Buffer
     
 end
 
-return { wrap = wrap }
\ No newline at end of file
+return { wrap = wrap, Buffer = Buffer, ScreenContent = ScreenContent }
\ No newline at end of file
diff --git a/lua/justfile b/lua/justfile
index 3056042..acf094b 100644
--- a/lua/justfile
+++ b/lua/justfile
@@ -1,7 +1,7 @@
 default:
     @just --list
 
-teal_files := "main.tl framebuffer.tl ringbuffer.tl"
+teal_files := "main.tl framebuffer.tl ringbuffer.tl socket.tl"
 
 build:
     mkdir -p out
diff --git a/lua/main.tl b/lua/main.tl
index 5fa7641..f26992f 100644
--- a/lua/main.tl
+++ b/lua/main.tl
@@ -1,74 +1,20 @@
 local json = require("json")
-local fb = require("framebuffer")
-local ringbuffer = require("ringbuffer")
+local Framebuffer = require("framebuffer")
+local Ringbuffer = require("ringbuffer")
+local Socket = require("socket")
 local UUID <const> = "8b9faf9f-9470-4a50-b405-0af5f0152550"
 local ENDPOINT <const> = "ws://localhost:8000/ipmi/computer/" .. UUID .. "/ws"
 
 print("[MAIN] Init")
 
-local enum SocketState
-    "reset"
-    "error"
-    "ok"
-    "viewer_connected"
-end
-
-local BAD_STATES <const> : {SocketState: boolean} = {
-    ["reset"] = true,
-    ["error"] = true,
-}
-
-local type SocketStateCallback = function(new_state: SocketState)
-
-local record Socket
-    state: SocketState
-    bad_state: function(self: Socket): boolean
-    set_state: function(self: Socket, state: SocketState)
-    on_state_change: function(self: Socket, cb: SocketStateCallback)
-    _callback: SocketStateCallback
-    ws: http.Websocket
-end
-
-local socket: Socket = {
-    state = "reset",
-    bad_state = function(self: Socket): boolean
-        return BAD_STATES[self.state] ~= nil
-    end,
-    set_state = function(self: Socket, state: SocketState)
-        self.state = state
-        self._callback(state)
-    end,
-    on_state_change = function(self: Socket, cb: SocketStateCallback)
-        self._callback = cb
-    end,
-    _callback = function(_: SocketState) end,
-    ws = nil,
-}
-
-local function send(message: string)
-    -- "message" needs to be valid JSON
-    -- otherwise the server will not accept it
-
-    if socket:bad_state() then return end
-
-    local r = { pcall(socket.ws.send, message) }
-    
-    if r[1] == false then
-        if (r[2] as string):sub(-11) == "closed file" then
-            socket:set_state("reset")
-        elseif (r[2] as string):sub(-9) == "too large" then
-            -- TODO handle
-            -- the connection stays open though
-        end
-    end
-end
+local socket = Socket.new(ENDPOINT)
 
 -- Set up framebuffer capture and statusline
 
 print("[MAIN] Setup framebuffer")
 
 local orig_native = term.native
-local buffer = fb.wrap(orig_native())
+local buffer = Framebuffer.wrap(orig_native())
 term.native = function(): term.Redirect
     return buffer.target
 end
@@ -89,30 +35,21 @@ end
 
 -- Create tasks
 
-local bar_codes: { SocketState: {string} } = {
+local bar_codes: { Socket.State: {string} } = {
     ["reset"] = {"[WS] RST", "78870111"},
     ["error"] = {"[WS] ERR", "78870EEE"},
     ["ok"] = {"[WS] OK\x03", "78870DD5"},
     ["viewer_connected"] = {"[WS] CON", "78870999"},
 }
 
-socket:on_state_change(function(new_state: SocketState)
+socket:on_state_change(function(new_state: Socket.State)
     set_bar(table.unpack(bar_codes[new_state]))
 end)
 
 local ws_task = coroutine.create(function()
     while true do
-        if socket:bad_state() then
-            --set_bar("[WS] RST", "78870111")
-            local r = http.websocket(ENDPOINT)
-            if r ~= false then
-                socket.ws = r as http.Websocket
-                --set_bar("[WS] OK\x03", "78870DD5")
-                socket:set_state("ok")
-            else
-                socket:set_state("error")
-                --set_bar("[WS] ERR", "78870EEE")
-            end
+        if socket:is_bad_state() then
+            socket:reconnect()
         end
         sleep(1)
     end
@@ -127,7 +64,7 @@ local report_task = coroutine.create(function()
             local message = json.encode({
                 screen = buffer.serialize()
             })
-            send(message)
+            socket:send(message)
             last_report = now
         end
         sleep(0) -- until next gametick
@@ -151,7 +88,7 @@ local tasks: {Task} = {
     {coro = report_task},
 }
 
-local event_queue = ringbuffer.new(64)
+local event_queue: Ringbuffer.Ringbuffer<table> = Ringbuffer.new(64)
 event_queue:push({n = 0})
 
 local shell_running = true
@@ -185,9 +122,7 @@ while shell_running do
     end
 end
 
-if socket.state == "ok" then
-    socket.ws.close()
-end
+socket:close()
 
 term.native = orig_native
 term.redirect(term.native())
diff --git a/lua/ringbuffer.tl b/lua/ringbuffer.tl
index ca7afe0..fb8c3e1 100644
--- a/lua/ringbuffer.tl
+++ b/lua/ringbuffer.tl
@@ -37,5 +37,6 @@ local function new<T>(size: integer): Ringbuffer<T>
 end
 
 return {
-    new = new
+    new = new,
+    Ringbuffer = Ringbuffer
 }
\ No newline at end of file
diff --git a/lua/socket.tl b/lua/socket.tl
new file mode 100644
index 0000000..7c86f21
--- /dev/null
+++ b/lua/socket.tl
@@ -0,0 +1,85 @@
+local enum State
+    "reset"
+    "error"
+    "ok"
+    "viewer_connected"
+end
+
+local BAD_STATES <const> : {State: boolean} = {
+    ["reset"] = true,
+    ["error"] = true,
+}
+
+local type StateCallback = function(new_state: State)
+
+local record Socket
+    state: State
+    is_bad_state: function(self: Socket): boolean
+    _set_state: function(self: Socket, state: State)
+    on_state_change: function(self: Socket, cb: StateCallback)
+    send: function(self: Socket, message: string)
+    reconnect: function(self: Socket)
+    close: function(self: Socket)
+    _endpoint: string
+    _callback: StateCallback
+    _ws: http.Websocket
+end
+
+local impl: table = {}
+
+impl.is_bad_state = function(self: Socket): boolean
+    return BAD_STATES[self.state] ~= nil
+end
+
+impl._set_state = function(self: Socket, state: State)
+    self.state = state
+    self._callback(state)
+end
+
+impl.on_state_change = function(self: Socket, cb: StateCallback)
+    self._callback = cb
+end
+
+impl.send = function(self: Socket, message: string)
+    -- "message" needs to be valid JSON
+    -- otherwise the server will not accept it
+
+    if self:is_bad_state() then return end
+
+    local r = { pcall(self._ws.send, message) }
+    
+    if r[1] == false then
+        if (r[2] as string):sub(-11) == "closed file" then
+            self:_set_state("reset")
+        elseif (r[2] as string):sub(-9) == "too large" then
+            -- TODO handle
+            -- the connection stays open though
+        end
+    end
+end
+
+impl.reconnect = function(self: Socket)
+    local r = http.websocket(self._endpoint)
+    if r ~= false then
+        self._ws = r as http.Websocket
+        self:_set_state("ok")
+    else
+        self:_set_state("error")
+    end
+end
+
+impl.close = function(self: Socket)
+    if self:is_bad_state() then return end
+    self._ws.close()
+end
+
+local function new(endpoint: string): Socket
+    return setmetatable({
+        state = "reset",
+        _endpoint = endpoint,
+        _callback = function(_: State) end,
+        _ws = nil,
+    }, { __index = impl })
+end
+
+return { new = new, State = State, StateCallback = StateCallback, Socket = Socket }
\ No newline at end of file
diff --git a/server/templates/install.lua b/server/templates/install.lua
index 6905f89..9f0d87b 100644
--- a/server/templates/install.lua
+++ b/server/templates/install.lua
@@ -1,5 +1,5 @@
 local path = "{{ deploy_path }}"
-files = { "main.lua", "json.lua", "framebuffer.lua", "ringbuffer.lua" }
+files = { "main.lua", "json.lua", "framebuffer.lua", "ringbuffer.lua", "socket.lua" }
 for _, file in ipairs(files) do
     fs.delete(file)
     shell.run(("wget %s/lua/%s"):format(path, file))