Merge pull request #6 from clams-tech/enhancement/restrict-origin

Enhancement: Restrict Origin

README.md

@@ -10,13 +10,19 @@ The code is inspired by and functions mostly the same as [ln-ws-proxy](https://g
 
 ## Updating Env Vars
 
-The WebSocket server is setup to run on port 3000, but can be modified by creating a .env file and placing it in the root directory with the following var:
+The following env vars will be picked up by the WebSocket server and can be set in your env to modify the defaults:
 
 ```
 HOST=localhost
 PORT=3000
 ```
 
+You can also add a `RESTRICT_ORIGINS` var if you would like the server to only accept connections from specified origins:
+
+```
+RESTRICT_ORIGINS=https://yourapp.com,https://staging.yourapp.com
+```
+
 ## Starting the server
 
 `yarn start`

package.json

@@ -1,6 +1,6 @@
 {
   "name": "lnsocket-proxy",
-  "version": "0.0.4",
+  "version": "0.0.5",
   "description": "A simple WebSocket server to proxy lnsocket frames to Core Lightning nodes over the Lightning Network",
   "main": "dist/index.js",
   "scripts": {

src/constants.ts

@@ -1,6 +1,6 @@
 // load env vars
 export const HOST = process.env.HOST || 'localhost'
-
 export const PORT = Number(process.env.PORT) || 3000
-
 export const MAX_SOCKET_BACKPRESSURE_BYTES = 1024 * 1024 // 1mb
+export const RESTRICT_ORIGINS =
+  process.env.RESTRICT_ORIGINS && process.env.RESTRICT_ORIGINS.split(',')

src/websockets/upgrade.ts

@@ -3,6 +3,7 @@ import { HttpResponse, HttpRequest, us_socket_context_t } from 'uWebSockets.js'
 import { RateLimiterMemory } from 'rate-limiter-flexible'
 import { arrayBufferToString, safetyPatchRes } from '../utils'
 import { Socket } from 'net'
+import { RESTRICT_ORIGINS } from '../constants'
 
 const connectionsRateLimiter = new RateLimiterMemory({
   points: 10, // connection attempts
@@ -26,6 +27,15 @@ async function handleUpgrade(
   const ip = arrayBufferToString(res.getRemoteAddressAsText())
   const nodeHost = req.getParameter(0)
 
+  if (RESTRICT_ORIGINS && !RESTRICT_ORIGINS.includes(origin)) {
+    res.cork(() => {
+      if (res.done) return
+      res.writeStatus('400 Bad Request').end()
+    })
+
+    return
+  }
+
   if (!nodeHost) {
     if (res.done) return