Added socket functions to lua environment

mod_lua.c

@@ -18,11 +18,112 @@ static void *mmap(void *addr, int64_t len, int prot, int flags, int fd,
 #include <sys/mman.h>
 #endif
 
-static void handle_request(struct mg_connection *);
+static const char *LUASOCKET = "luasocket";
 
+// Forward declarations
+static void handle_request(struct mg_connection *);
 static int handle_lsp_request(struct mg_connection *, const char *,
                               struct file *, struct lua_State *);
 
+static void reg_string(struct lua_State *L, const char *name, const char *val) {
+  lua_pushstring(L, name);
+  lua_pushstring(L, val);
+  lua_rawset(L, -3);
+}
+
+static void reg_int(struct lua_State *L, const char *name, int val) {
+  lua_pushstring(L, name);
+  lua_pushinteger(L, val);
+  lua_rawset(L, -3);
+}
+
+static void reg_function(struct lua_State *L, const char *name,
+                         lua_CFunction func, struct mg_connection *conn) {
+  lua_pushstring(L, name);
+  lua_pushlightuserdata(L, conn);
+  lua_pushcclosure(L, func, 1);
+  lua_rawset(L, -3);
+}
+
+static int lsp_sock_close(lua_State *L) {
+  if (lua_gettop(L) > 0 && lua_istable(L, -1)) {
+    lua_getfield(L, -1, "sock");
+    closesocket((SOCKET) lua_tonumber(L, -1));
+  } else {
+    return luaL_error(L, "invalid :close() call");
+  }
+  return 1;
+}
+
+static int lsp_sock_recv(lua_State *L) {
+  char buf[2000];
+  int n;
+
+  if (lua_gettop(L) > 0 && lua_istable(L, -1)) {
+    lua_getfield(L, -1, "sock");
+    n = recv((SOCKET) lua_tonumber(L, -1), buf, sizeof(buf), 0);
+    if (n <= 0) {
+      lua_pushnil(L);
+    } else {
+      lua_pushlstring(L, buf, n);
+    }
+  } else {
+    return luaL_error(L, "invalid :close() call");
+  }
+  return 1;
+}
+
+static int lsp_sock_send(lua_State *L) {
+  const char *buf;
+  size_t len, sent = 0;
+  int n, sock;
+
+  if (lua_gettop(L) > 1 && lua_istable(L, -2) && lua_isstring(L, -1)) {
+    buf = lua_tolstring(L, -1, &len);
+    lua_getfield(L, -2, "sock");
+    sock = lua_tonumber(L, -1);
+    while (sent < len) {
+      if ((n = send(sock, buf + sent, len - sent, 0)) <= 0) {
+        break;
+      }
+      sent += n;
+    }
+    lua_pushnumber(L, n);
+  } else {
+    return luaL_error(L, "invalid :close() call");
+  }
+  return 1;
+}
+
+static const struct luaL_Reg luasocket_methods[] = {
+  {"close", lsp_sock_close},
+  {"send", lsp_sock_send},
+  {"recv", lsp_sock_recv},
+  {NULL, NULL}
+};
+
+static int lsp_connect(lua_State *L) {
+  char ebuf[100];
+  SOCKET sock;
+
+  if (lua_isstring(L, -3) && lua_isnumber(L, -2) && lua_isnumber(L, -1)) {
+    sock = conn2(lua_tostring(L, -3), lua_tonumber(L, -2), lua_tonumber(L, -1),
+                 ebuf, sizeof(ebuf));
+    if (sock == INVALID_SOCKET) {
+      return luaL_error(L, ebuf);
+    } else {
+      lua_newtable(L);
+      reg_int(L, "sock", sock);
+      reg_string(L, "host", lua_tostring(L, -4));
+      luaL_getmetatable(L, LUASOCKET);
+      lua_setmetatable(L, -2);
+    }
+  } else {
+    return luaL_error(L, "connect(host,port,is_ssl): invalid parameter given.");
+  }
+  return 1;
+}
+
 static int lsp_error(lua_State *L) {
   lua_getglobal(L, "mg");
   lua_getfield(L, -1, "onerror");
@@ -141,26 +242,6 @@ static int lsp_redirect(lua_State *L) {
   return 0;
 }
 
-static void reg_string(struct lua_State *L, const char *name, const char *val) {
-  lua_pushstring(L, name);
-  lua_pushstring(L, val);
-  lua_rawset(L, -3);
-}
-
-static void reg_int(struct lua_State *L, const char *name, int val) {
-  lua_pushstring(L, name);
-  lua_pushinteger(L, val);
-  lua_rawset(L, -3);
-}
-
-static void reg_function(struct lua_State *L, const char *name,
-                         lua_CFunction func, struct mg_connection *conn) {
-  lua_pushstring(L, name);
-  lua_pushlightuserdata(L, conn);
-  lua_pushcclosure(L, func, 1);
-  lua_rawset(L, -3);
-}
-
 void mg_prepare_lua_environment(struct mg_connection *conn, lua_State *L) {
   const struct mg_request_info *ri = mg_get_request_info(conn);
   extern void luaL_openlibs(lua_State *);
@@ -171,6 +252,13 @@ void mg_prepare_lua_environment(struct mg_connection *conn, lua_State *L) {
   { extern int luaopen_lsqlite3(lua_State *); luaopen_lsqlite3(L); }
 #endif
 
+  luaL_newmetatable(L, LUASOCKET);
+  lua_pushliteral(L, "__index");
+  luaL_newlib(L, luasocket_methods);
+  lua_rawset(L, -3);
+  lua_pop(L, 1);
+  lua_register(L, "connect", lsp_connect);
+
   if (conn == NULL) return;
 
   // Register mg module

mongoose.c

@@ -514,6 +514,13 @@ struct mg_connection {
   int64_t last_throttle_bytes;// Bytes sent this second
 };
 
+// Directory entry
+struct de {
+  struct mg_connection *conn;
+  char *file_name;
+  struct file file;
+};
+
 const char **mg_get_valid_option_names(void) {
   return config_options;
 }
@@ -2508,11 +2515,36 @@ int mg_modify_passwords_file(const char *fname, const char *domain,
   return 1;
 }
 
-struct de {
-  struct mg_connection *conn;
-  char *file_name;
-  struct file file;
-};
+static int conn2(const char *host, int port, int use_ssl,
+                 char *ebuf, size_t ebuf_len) {
+  struct sockaddr_in sin;
+  struct hostent *he;
+  SOCKET sock = INVALID_SOCKET;
+
+  if (host == NULL) {
+    snprintf(ebuf, ebuf_len, "%s", "NULL host");
+  } else if (use_ssl && SSLv23_client_method == NULL) {
+    snprintf(ebuf, ebuf_len, "%s", "SSL is not initialized");
+    // TODO(lsm): use something threadsafe instead of gethostbyname()
+  } else if ((he = gethostbyname(host)) == NULL) {
+    snprintf(ebuf, ebuf_len, "gethostbyname(%s): %s", host, strerror(ERRNO));
+  } else if ((sock = socket(PF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) {
+    snprintf(ebuf, ebuf_len, "socket(): %s", strerror(ERRNO));
+  } else {
+    sin.sin_family = AF_INET;
+    sin.sin_port = htons((uint16_t) port);
+    sin.sin_addr = * (struct in_addr *) he->h_addr_list[0];
+    if (connect(sock, (struct sockaddr *) &sin, sizeof(sin)) != 0) {
+      snprintf(ebuf, ebuf_len, "connect(%s:%d): %s",
+               host, port, strerror(ERRNO));
+      closesocket(sock);
+      sock = INVALID_SOCKET;
+    }
+  }
+  return sock;
+}
+
+
 
 static void url_encode(const char *src, char *dst, size_t dst_len) {
   static const char *dont_escape = "._-$,;~()";
@@ -4641,54 +4673,37 @@ struct mg_connection *mg_connect(const char *host, int port, int use_ssl,
                                  char *ebuf, size_t ebuf_len) {
   static struct mg_context fake_ctx;
   struct mg_connection *conn = NULL;
-  struct sockaddr_in sin;
-  struct hostent *he;
   SOCKET sock;
 
-  if (host == NULL) {
-    snprintf(ebuf, ebuf_len, "%s", "NULL host");
-  } else if (use_ssl && SSLv23_client_method == NULL) {
-    snprintf(ebuf, ebuf_len, "%s", "SSL is not initialized");
-  } else if ((he = gethostbyname(host)) == NULL) {
-    snprintf(ebuf, ebuf_len, "gethostbyname(%s): %s", host, strerror(ERRNO));
-  } else if ((sock = socket(PF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) {
-    snprintf(ebuf, ebuf_len, "socket(): %s", strerror(ERRNO));
-  } else {
-    sin.sin_family = AF_INET;
-    sin.sin_port = htons((uint16_t) port);
-    sin.sin_addr = * (struct in_addr *) he->h_addr_list[0];
-    if (connect(sock, (struct sockaddr *) &sin, sizeof(sin)) != 0) {
-      snprintf(ebuf, ebuf_len, "connect(%s:%d): %s",
-               host, port, strerror(ERRNO));
-      closesocket(sock);
-    } else if ((conn = (struct mg_connection *)
-                calloc(1, sizeof(*conn) + MAX_REQUEST_SIZE)) == NULL) {
-      snprintf(ebuf, ebuf_len, "calloc(): %s", strerror(ERRNO));
-      closesocket(sock);
+  if ((sock = conn2(host, port, use_ssl, ebuf, ebuf_len)) == INVALID_SOCKET) {
+  } else if ((conn = (struct mg_connection *)
+              calloc(1, sizeof(*conn) + MAX_REQUEST_SIZE)) == NULL) {
+    snprintf(ebuf, ebuf_len, "calloc(): %s", strerror(ERRNO));
+    closesocket(sock);
 #ifndef NO_SSL
-    } else if (use_ssl && (conn->client_ssl_ctx =
-                           SSL_CTX_new(SSLv23_client_method())) == NULL) {
-      snprintf(ebuf, ebuf_len, "SSL_CTX_new error");
-      closesocket(sock);
-      free(conn);
-      conn = NULL;
+  } else if (use_ssl && (conn->client_ssl_ctx =
+                         SSL_CTX_new(SSLv23_client_method())) == NULL) {
+    snprintf(ebuf, ebuf_len, "SSL_CTX_new error");
+    closesocket(sock);
+    free(conn);
+    conn = NULL;
 #endif // NO_SSL
-    } else {
-      conn->buf_size = MAX_REQUEST_SIZE;
-      conn->buf = (char *) (conn + 1);
-      conn->ctx = &fake_ctx;
-      conn->client.sock = sock;
-      conn->client.rsa.sin = sin;
-      conn->client.is_ssl = use_ssl;
+  } else {
+    socklen_t len;
+    conn->buf_size = MAX_REQUEST_SIZE;
+    conn->buf = (char *) (conn + 1);
+    conn->ctx = &fake_ctx;
+    conn->client.sock = sock;
+    getsockname(sock, &conn->client.rsa.sa, &len);
+    conn->client.is_ssl = use_ssl;
 #ifndef NO_SSL
-      if (use_ssl) {
-        // SSL_CTX_set_verify call is needed to switch off server certificate
-        // checking, which is off by default in OpenSSL and on in yaSSL.
-        SSL_CTX_set_verify(conn->client_ssl_ctx, 0, 0);
-        sslize(conn, conn->client_ssl_ctx, SSL_connect);
-      }
-#endif
+    if (use_ssl) {
+      // SSL_CTX_set_verify call is needed to switch off server certificate
+      // checking, which is off by default in OpenSSL and on in yaSSL.
+      SSL_CTX_set_verify(conn->client_ssl_ctx, 0, 0);
+      sslize(conn, conn->client_ssl_ctx, SSL_connect);
     }
+#endif
   }
 
   return conn;