Adding in ssl certification

codalab/bin/ws_server.py

@@ -2,12 +2,13 @@
 import argparse
 import asyncio
 from collections import defaultdict
+from dataclasses import dataclass
 import logging
 import os
 import re
 from typing import Any, Dict, Optional
+import ssl
 import websockets
-from dataclasses import dataclass
 import threading
 
 from codalab.lib.codalab_manager import CodaLabManager
@@ -52,16 +53,20 @@ def last_use(self):
     def last_use(self, value):
         self._last_use = value
 
+# Configure logging.
+logger = logging.getLogger(__name__)
 
-worker_to_ws: Dict[str, Dict[str, WS]] = defaultdict(
-    dict
-)  # Maps worker to list of its websockets (since each worker has a pool of connections)
+# Variables for authenticating and acknowledging receipt of messages.
+server_secret = os.environ["CODALAB_SERVER_SECRET"]
 ACK = b'a'
-logger = logging.getLogger(__name__)
+
+# Maps worker to list of its websockets (since each worker has a pool of connections)
+worker_to_ws: Dict[str, Dict[str, WS]] = defaultdict(dict)  
+
+# Configure access to SQL database.
 manager = CodaLabManager()
 bundle_model = manager.model()
 worker_model = manager.worker_model()
-server_secret = os.environ["CODALAB_SERVER_SECRET"]
 
 
 async def send_handler(server_websocket, worker_id):
@@ -135,9 +140,20 @@ async def async_main():
     parser.add_argument(
         '--port', help='Port to run the server on.', type=int, required=False, default=2901
     )
+    parser.add_argument(
+        '--cert-path', help='Path to SSL certificate file for websocket server.', type=str, required=False
+    )
     args = parser.parse_args()
+
+    # Parse SSL config.
+    if args.cert_path:
+        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        ssl_context.load_cert_chain(args.cert_path)
+    else:
+        ssl_context = None
+
     logging.debug(f"Running ws-server on 0.0.0.0:{args.port}")
-    async with websockets.serve(ws_handler, "0.0.0.0", args.port):
+    async with websockets.serve(ws_handler, "0.0.0.0", args.port, ssl=ssl_context):
         await asyncio.Future()  # run server forever
 
 

codalab/model/worker_model.py

@@ -474,8 +474,8 @@ def send_json(self, data: dict, worker_id: str, timeout_secs: int = 60):
         """
         start_time = time.time()
         while time.time() - start_time < timeout_secs:
-            try:
-                with connect(f"{self._ws_server}/send/{worker_id}") as websocket:
+            try:  # TODO(agaut): don't always set ssl to be True
+                with connect(f"{self._ws_server}/send/{worker_id}", ssl=True) as websocket:
                     websocket.send(self._server_secret)  # Authenticate
                     websocket.send(json.dumps(data).encode())
                     ack = websocket.recv()

codalab/worker/worker.py

@@ -359,8 +359,8 @@ async def listen(self, thread_id):
         wss_uri = f"{self.ws_server}/worker/{self.id}/{thread_id}"
         while not self.terminate:
             logger.info(f"Connecting to {wss_uri}")
-            try:
-                async with websockets.connect(f"{wss_uri}", max_queue=1) as websocket:
+            try:  # TODO(agaut): don't always set ssl to be True
+                async with websockets.connect(f"{wss_uri}", max_queue=1, ssl=True) as websocket:
                     await self.recv_messages(websocket)
             except Exception:
                 logger.error(f"Error connecting to ws-server: {traceback.print_exc()}")

docker_config/compose_files/docker-compose.ssl.yml

@@ -8,3 +8,7 @@ services:
       - ./files/nginx.conf.ssl:/etc/nginx/nginx.conf:ro
       - ${CODALAB_SSL_KEY_FILE}:/opt/ssl/codalab.key
       - ${CODALAB_SSL_CERT_FILE}:/opt/ssl/codalab.crt
+  ws-server:
+    volumes:
+      - ${CODALAB_SSL_KEY_FILE}:/opt/ssl/codalab.key
+      - ${CODALAB_SSL_CERT_FILE}:/opt/ssl/codalab.crt

docker_config/compose_files/docker-compose.yml

@@ -133,7 +133,7 @@ services:
 
   ws-server:
     image: codalab/server:${CODALAB_VERSION}
-    command: cl-ws-server --port ${CODALAB_WS_PORT}
+    command: cl-ws-server --port ${CODALAB_WS_PORT} --cert-path ${CODALAB_SSL_CERT_FILE}
     <<: *codalab-base
     <<: *codalab-server
     depends_on: