diff --git a/src/index.js b/src/index.js
index 26de44a..7a6cafa 100644
--- a/src/index.js
+++ b/src/index.js
@@ -5,7 +5,7 @@ const mafmt = require('mafmt')
 const withIs = require('class-is')
 const toUri = require('multiaddr-to-uri')
 const log = require('debug')('libp2p:websockets:transport')
-const abortable = require('abortable-iterator')
+const Libp2pSocket = require('./socket')
 const { AbortError } = require('interface-transport')
 const createListener = require('./listener')
 
@@ -44,20 +44,7 @@ class WebSockets {
     }
 
     log('connected %s', ma)
-    return {
-      sink: async source => {
-        try {
-          await socket.sink(abortable(source, options.signal))
-        } catch (err) {
-          // Re-throw non-aborted errors
-          if (err.type !== 'aborted') throw err
-          // Otherwise, this is fine...
-          await socket.close()
-        }
-      },
-      source: abortable(socket.source, options.signal),
-      getObservedAddrs
-    }
+    return new Libp2pSocket(socket, ma, options)
   }
 
   createListener (options, handler) {
diff --git a/src/socket.js b/src/socket.js
new file mode 100644
index 0000000..4a88862
--- /dev/null
+++ b/src/socket.js
@@ -0,0 +1,32 @@
+'use strict'
+
+const abortable = require('abortable-iterator')
+
+const { Connection } = require('interface-connection')
+
+class Libp2pSocket extends Connection {
+  constructor (rawSocket, ma, opts = {}) {
+    super(ma)
+
+    this._rawSocket = rawSocket
+    this._ma = ma
+
+    this.sink = this._sink(opts)
+    this.source = opts.signal ? abortable(rawSocket.source, opts.signal) : rawSocket.source
+  }
+
+  _sink (opts) {
+    return async (source) => {
+      try {
+        await this._rawSocket.sink(abortable(source, opts.signal))
+      } catch (err) {
+        // Re-throw non-aborted errors
+        if (err.type !== 'aborted') throw err
+        // Otherwise, this is fine...
+        await this._rawSocket.close()
+      }
+    }
+  }
+}
+
+module.exports = Libp2pSocket