From 22461d17281d9095f0cf2f5f689c4c0796b6a197 Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Sat, 13 Feb 2021 08:09:44 -0600
Subject: [PATCH] simple echo websocket working

---
 docker/Dockerfile.base |  5 ++++-
 frigate/app.py         | 10 +++++++---
 frigate/http.py        | 10 ++++++++++
 nginx/nginx.conf       |  8 ++++++++
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base
index 1acbe4c46..794e81dac 100644
--- a/docker/Dockerfile.base
+++ b/docker/Dockerfile.base
@@ -34,7 +34,10 @@ RUN apt-get -qq update \
 RUN pip3 install \
     peewee_migrate \
     zeroconf \
-    voluptuous
+    voluptuous\
+    Flask-Sockets \
+    gevent \
+    gevent-websocket
 
 COPY nginx/nginx.conf /etc/nginx/nginx.conf
 
diff --git a/frigate/app.py b/frigate/app.py
index f4032bc59..e3e63802f 100644
--- a/frigate/app.py
+++ b/frigate/app.py
@@ -8,6 +8,8 @@ import sys
 import signal
 
 import yaml
+from gevent import pywsgi
+from geventwebsocket.handler import WebSocketHandler
 from peewee_migrate import Router
 from playhouse.sqlite_ext import SqliteExtDatabase
 from playhouse.sqliteq import SqliteQueueDatabase
@@ -106,8 +108,8 @@ class FrigateApp():
         for log, level in self.config.logger.logs.items():
             logging.getLogger(log).setLevel(level)
         
-        if not 'werkzeug' in self.config.logger.logs:
-            logging.getLogger('werkzeug').setLevel('ERROR')
+        if not 'geventwebsocket.handler' in self.config.logger.logs:
+            logging.getLogger('geventwebsocket.handler').setLevel('ERROR')
 
     def init_queues(self):
         # Queues for clip processing
@@ -239,7 +241,9 @@ class FrigateApp():
         
         signal.signal(signal.SIGTERM, receiveSignal)
 
-        self.flask_app.run(host='127.0.0.1', port=5001, debug=False)
+        server = pywsgi.WSGIServer(('127.0.0.1', 5001), self.flask_app, handler_class=WebSocketHandler)
+        server.serve_forever()
+
         self.stop()
     
     def stop(self):
diff --git a/frigate/http.py b/frigate/http.py
index e6a87ecc5..d292281a6 100644
--- a/frigate/http.py
+++ b/frigate/http.py
@@ -9,6 +9,7 @@ import cv2
 import numpy as np
 from flask import (Blueprint, Flask, Response, current_app, jsonify,
                    make_response, request)
+from flask_sockets import Sockets
 from peewee import SqliteDatabase, operator, fn, DoesNotExist
 from playhouse.shortcuts import model_to_dict
 
@@ -21,9 +22,11 @@ from frigate.version import VERSION
 logger = logging.getLogger(__name__)
 
 bp = Blueprint('frigate', __name__)
+ws = Blueprint('ws', __name__)
 
 def create_app(frigate_config, database: SqliteDatabase, stats_tracking, detected_frames_processor):
     app = Flask(__name__)
+    sockets = Sockets(app)
 
     @app.before_request
     def _db_connect():
@@ -39,6 +42,7 @@ def create_app(frigate_config, database: SqliteDatabase, stats_tracking, detecte
     app.detected_frames_processor = detected_frames_processor
 
     app.register_blueprint(bp)
+    sockets.register_blueprint(ws)
 
     return app
 
@@ -304,3 +308,9 @@ def imagestream(detected_frames_processor, camera_name, fps, height, draw_option
         ret, jpg = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
         yield (b'--frame\r\n'
             b'Content-Type: image/jpeg\r\n\r\n' + jpg.tobytes() + b'\r\n\r\n')
+
+@ws.route('/ws')
+def echo_socket(socket):
+    while not socket.closed:
+        message = socket.receive()
+        socket.send(message)
diff --git a/nginx/nginx.conf b/nginx/nginx.conf
index 4c0b4638f..51842cb70 100644
--- a/nginx/nginx.conf
+++ b/nginx/nginx.conf
@@ -102,6 +102,14 @@ http {
             root /media/frigate;
         }
 
+        location /ws {
+            proxy_pass http://frigate_api/ws;
+            proxy_http_version 1.1;
+            proxy_set_header Upgrade $http_upgrade;
+            proxy_set_header Connection "Upgrade";
+            proxy_set_header Host $host;
+        }
+
         location /api/ {
             add_header 'Access-Control-Allow-Origin' '*';
             add_header Cache-Control "no-store";