From: Shuanglei Tao <tsl0922@gmail.com>
Date: Sat, 24 Sep 2016 16:56:53 +0000 (+0800)
Subject: Implemented authentication for websocket connection
X-Git-Url: http://git.prime8.dev/?a=commitdiff_plain;h=9174d6116cd6448cda0ae237fc124473c9add12d;p=ttyd.git

Implemented authentication for websocket connection
---

diff --git a/src/http.c b/src/http.c
index 7147ec5..52e207f 100644
--- a/src/http.c
+++ b/src/http.c
@@ -18,11 +18,11 @@ check_auth(struct lws *wsi) {
             if (strlen(token) == 0)
                 continue;
             if (i++ == 2) {
-                b64_text = strdup(token);
+                b64_text = token;
                 break;
             }
         }
-        if (b64_text != NULL && strcmp(b64_text, server->credential) == 0)
+        if (b64_text != NULL && !strcmp(b64_text, server->credential))
             return 0;
     }
 
@@ -51,18 +51,13 @@ int
 callback_http(struct lws *wsi, enum lws_callback_reasons reason, void *user, void *in, size_t len) {
     unsigned char buffer[4096 + LWS_PRE], *p, *end;
     char buf[256];
-    int n;
 
     switch (reason) {
         case LWS_CALLBACK_HTTP:
-            lwsl_notice("lws_http_serve: %s\n", in);
-
             {
                 char name[100], rip[50];
-                lws_get_peer_addresses(wsi, lws_get_socket_fd(wsi), name,
-                                       sizeof(name), rip, sizeof(rip));
-                sprintf(buf, "%s (%s)", name, rip);
-                lwsl_notice("HTTP connect from %s\n", buf);
+                lws_get_peer_addresses(wsi, lws_get_socket_fd(wsi), name, sizeof(name), rip, sizeof(rip));
+                lwsl_notice("HTTP connect from %s (%s), path: %s\n", name, rip, in);
             }
 
             if (len < 1) {
@@ -70,7 +65,6 @@ callback_http(struct lws *wsi, enum lws_callback_reasons reason, void *user, voi
                 goto try_to_reuse;
             }
 
-            // TODO: this doesn't work for websocket
             switch (check_auth(wsi)) {
                 case 0:
                     break;
@@ -85,37 +79,54 @@ callback_http(struct lws *wsi, enum lws_callback_reasons reason, void *user, voi
             if (lws_hdr_total_length(wsi, WSI_TOKEN_POST_URI))
                 return 0;
 
-            if (strcmp((const char *) in, "/")) {
-                lws_return_http_status(wsi, HTTP_STATUS_NOT_FOUND, NULL);
+            p = buffer + LWS_PRE;
+            end = p + sizeof(buffer) - LWS_PRE;
+
+            if (!strncmp((const char *)in, "/auth_token.js", 14)) {
+                size_t n = server->credential != NULL ? sprintf(buf, "var tty_auth_token = '%s';", server->credential) : 0;
+
+                if (lws_add_http_header_status(wsi, HTTP_STATUS_OK, &p, end))
+                    return 1;
+                if (lws_add_http_header_by_token(wsi,
+                                                 WSI_TOKEN_HTTP_CONTENT_TYPE,
+                                                 (unsigned char *) "application/javascript",
+                                                 22, &p, end))
+                    return 1;
+                if (lws_add_http_header_content_length(wsi, (unsigned long) n, &p, end))
+                    return 1;
+                if (lws_finalize_http_header(wsi, &p, end))
+                    return 1;
+                if (lws_write(wsi, buffer + LWS_PRE, p - (buffer + LWS_PRE), LWS_WRITE_HTTP_HEADERS) < 0)
+                    return 1;
+                if (n > 0 && lws_write_http(wsi, buf, n) < 0) {
+                    return 1;
+                }
                 goto try_to_reuse;
             }
 
-            p = buffer + LWS_PRE;
-            end = p + sizeof(buffer) - LWS_PRE;
+            if (strncmp((const char *) in, "/", 1)) {
+                lws_return_http_status(wsi, HTTP_STATUS_NOT_FOUND, NULL);
+                goto try_to_reuse;
+            }
 
-            if (lws_add_http_header_status(wsi, 200, &p, end))
+            if (lws_add_http_header_status(wsi, HTTP_STATUS_OK, &p, end))
                 return 1;
             if (lws_add_http_header_by_token(wsi,
                                              WSI_TOKEN_HTTP_CONTENT_TYPE,
                                              (unsigned char *) "text/html",
                                              9, &p, end))
                 return 1;
-            if (lws_add_http_header_content_length(wsi, index_html_len, &p, end))
+            if (lws_add_http_header_content_length(wsi, (unsigned long) index_html_len, &p, end))
                 return 1;
             if (lws_finalize_http_header(wsi, &p, end))
                 return 1;
-            n = lws_write(wsi, buffer + LWS_PRE, p - (buffer + LWS_PRE), LWS_WRITE_HTTP_HEADERS);
-            if (n < 0) {
+            if (lws_write(wsi, buffer + LWS_PRE, p - (buffer + LWS_PRE), LWS_WRITE_HTTP_HEADERS) < 0) {
                 return 1;
             }
 
-            n = lws_write_http(wsi, index_html, index_html_len);
-            if (n < 0)
+            if (lws_write_http(wsi, index_html, index_html_len) < 0)
                 return 1;
             goto try_to_reuse;
-        case LWS_CALLBACK_HTTP_WRITEABLE:
-            lwsl_info("LWS_CALLBACK_HTTP_WRITEABLE\n");
-            break;
         default:
             break;
     }
diff --git a/src/index.html b/src/index.html
index 37d54d6..8a15414 100644
--- a/src/index.html
+++ b/src/index.html
@@ -571,6 +571,7 @@
             "Ch6PdwAAAGyWjFW4yScjaWa2mGcofHxWxewKALglWBpLUvwwk+UOh5eNGyUOs1/EF+pZr+ud5Ozo"+"GwYdAABg2p52LiSgAY/ZVlOmilEgHn6G3OcwYjzI7vOj1t6xsx4S3lBY96EUQBF6AIBAmPYH4PoG"+"YCoJAADWe+OZJZi7/x76/yH7Lzf9M5XzRKnFPmveMsilQHwVAAAAAKB3LQD8PCIAAADga0QujBLy"+"wzeJ4a6Z/ERVBAUlAEDqvoM7BQBAuAguzFqILtmjH3Kd4wfKobnOhA3z85qWoRPm9hwoOHoDAAlC"+"bwDAA56FHAuXflHo3fe2ttG9XUDeA9YmYCBQ0oPr/1QC8IvuCwAAApbUAQCK22MmE3O78VAbHQT9"+"PIPNoT9zNc3l2Oe7TAVLANBufT8MAQAAAGzT4PS8AQAAoELGHb2uaCwwEv1EWhFriUkbAaAZ27/f"+"VZnTZXbWz3BwWpjUaMZKRj7dZ0J//gUeTdpVEwAAZOFsNxKAjQSgA+ABPoY8Jj5y2wje81jsXc/1"+
             "TOQWTDYZBmAkNDiqVwuA2NJ9AQAAEBKAt9Vrsfs/2N19MO91S9rd8EHTZHnzC5MYmfQEACy/FBcA"+"AADA5c4gi4z8RANs/m6FNXVo9DV46JG1BBDukqlw/Va5G7QbuGVSI+2aZaoLXJrdVj2zlC9Z5QEA"+"EFz/5QzgVZwAAAAA/oXcxyC6WfTu+09Ve/c766J4VTAGUFmA51+VANKi/QPoPwYgYAkA715OH4S0"+"s5KDHvj99MMq8TPFc3roKZnGOoT1bmIhVgc7XAMBAAAAAMAW1VbQw3gapzOpJd+Kd2fc4iSO62fJ"+"v9+movui1wUNPAj059N3OVxzk4gV73PmE8FIA2F5mRq37Evc76vLXfF4rD5UJJAw46hW6LZCb5sN"+"Ldx+kzMCAAB+hfy95+965ZCLP7B3/VlTHCvDEKtQhTm4KiCgAEAbrfbWTPssAAAAXpee1tVrozYY"+"n41wD1aeYtkKfswN5/SXPO0JDnhO/4laUortv/s412fybe/nONdncoCHnBVliu0CQGBWlPY/5Kwo"+
             "m2L/kruPM6Q7oz4tvDQy+bZ3HzOi+gNHA4DZEgA="+"");lib.resource.add("hterm/concat/date","text/plain","Sat, 10 Sep 2016 08:51:57 +0000"+"");lib.resource.add("hterm/changelog/version","text/plain","1.58"+"");lib.resource.add("hterm/changelog/date","text/plain","2016-07-12"+"");lib.resource.add("hterm/git/HEAD","text/plain","49f8641dd055afaad9eadcd8553804eff0dd2637"+"");</script>
+    <script src="auth_token.js"></script>
     <script type="text/javascript">
         (function() {
             var httpsEnabled = window.location.protocol == "https:";
@@ -580,12 +581,12 @@
 
             var openWs = function() {
                 var ws = new WebSocket(url, protocols);
-
-                var term;
-
-                var pingTimer;
+                var term, pingTimer;
 
                 ws.onopen = function(event) {
+                    if (typeof tty_auth_token !== 'undefined') {
+                        ws.send(JSON.stringify({AuthToken: tty_auth_token}));
+                    }
                     pingTimer = setInterval(sendPing, 30 * 1000, ws);
 
                     hterm.defaultStorage = new lib.Storage.Local();
diff --git a/src/protocol.c b/src/protocol.c
index b581e27..680a77f 100644
--- a/src/protocol.c
+++ b/src/protocol.c
@@ -4,6 +4,7 @@
 #define INPUT '0'
 #define PING '1'
 #define RESIZE_TERMINAL '2'
+#define JSON_DATA '{'
 
 // server message
 #define OUTPUT '0'
@@ -158,6 +159,7 @@ callback_tty(struct lws *wsi, enum lws_callback_reasons reason,
         case LWS_CALLBACK_ESTABLISHED:
             client->exit = false;
             client->initialized = false;
+            client->authenticated = false;
             client->wsi = wsi;
             lws_get_peer_addresses(wsi, lws_get_socket_fd(wsi),
                                    client->hostname, sizeof(client->hostname),
@@ -165,7 +167,7 @@ callback_tty(struct lws *wsi, enum lws_callback_reasons reason,
             STAILQ_INIT(&client->queue);
             if (pthread_create(&client->thread, NULL, thread_run_command, client) != 0) {
                 lwsl_err("pthread_create\n");
-                return 1;
+                return -1;
             }
 
             pthread_mutex_lock(&server->lock);
@@ -222,6 +224,13 @@ callback_tty(struct lws *wsi, enum lws_callback_reasons reason,
         case LWS_CALLBACK_RECEIVE:
             data = (char *) in;
             char command = data[0];
+
+            // check auth
+            if (server->credential != NULL && !client->authenticated && command != JSON_DATA) {
+                lwsl_notice("websocket authentication failed\n");
+                return -1;
+            }
+
             switch (command) {
                 case INPUT:
                     if (write(client->pty, data + 1, len - 1) < len - 1) {
@@ -247,6 +256,22 @@ callback_tty(struct lws *wsi, enum lws_callback_reasons reason,
                         t_free(size);
                     }
                     break;
+                case JSON_DATA:
+                    if (server->credential == NULL)
+                        break;
+                    {
+                        json_object *obj = json_tokener_parse(data);
+                        struct json_object *o = NULL;
+                        if (json_object_object_get_ex(obj, "AuthToken", &o)) {
+                            const char *token = json_object_get_string(o);
+                            if (strcmp(token, server->credential)) {
+                                lwsl_notice("websocket authentication failed with token: %s\n", token);
+                                return -1;
+                            }
+                        }
+                        client->authenticated = true;
+                    }
+                    break;
                 default:
                     lwsl_notice("unknown message type: %c\n", command);
                     break;
diff --git a/src/server.h b/src/server.h
index a71264a..ae4bad4 100644
--- a/src/server.h
+++ b/src/server.h
@@ -18,9 +18,7 @@
 #include <assert.h>
 
 #ifdef __APPLE__
-
 #include <util.h>
-
 #else
 #include <pty.h>
 #endif
@@ -43,6 +41,7 @@ struct pty_data {
 struct tty_client {
     bool exit;
     bool initialized;
+    bool authenticated;
     char hostname[100];
     char address[50];
 
diff --git a/src/utils.c b/src/utils.c
index f32564e..e8e2430 100644
--- a/src/utils.c
+++ b/src/utils.c
@@ -36,8 +36,7 @@ t_malloc(size_t size) {
 }
 
 void t_free(void *p) {
-    if (p)
-        free(p);
+    free(p);
 }
 
 void *