From 0e669f6c8203ea7fcb49d40f7e3474c405013e9d Mon Sep 17 00:00:00 2001
From: Bruno Silvestre <brunoos@inf.ufg.br>
Date: Sun, 19 Feb 2023 08:56:24 -0300
Subject: [PATCH] Add identity hint to PSK

---
 samples/psk/server.lua |  8 ++++++++
 src/context.c          | 13 +++++++++++++
 src/ssl.lua            | 26 +++++++++++++++++++-------
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/samples/psk/server.lua b/samples/psk/server.lua
index 245c5cd..b5a958b 100644
--- a/samples/psk/server.lua
+++ b/samples/psk/server.lua
@@ -19,7 +19,15 @@ local params = {
    mode = "server",
    protocol = "any",
    options = "all",
+
+-- PSK with just a callback
    psk = pskcb,
+
+-- PSK with identity hint
+--   psk = {
+--      hint = "hintpsksample",
+--      callback = pskcb,
+--   },
 }
 
 
diff --git a/src/context.c b/src/context.c
index 02198a8..dcabbc0 100644
--- a/src/context.c
+++ b/src/context.c
@@ -764,6 +764,18 @@ static int set_server_psk_cb(lua_State *L)
   return 1;
 }
 
+/*
+ * Set the PSK indentity hint.
+ */
+static int set_psk_identity_hint(lua_State *L)
+{
+  p_context ctx = checkctx(L, 1);
+  const char *hint = luaL_checkstring(L, 2);
+  int ret = SSL_CTX_use_psk_identity_hint(ctx->context, hint);
+  lua_pushboolean(L, ret);
+  return 1;
+}
+
 /*
  * Client callback to PSK.
  */
@@ -882,6 +894,7 @@ static luaL_Reg funcs[] = {
   {"setdhparam",      set_dhparam},
   {"setverify",       set_verify},
   {"setoptions",      set_options},
+  {"setpskhint",      set_psk_identity_hint},
   {"setserverpskcb",  set_server_psk_cb},
   {"setclientpskcb",  set_client_psk_cb},
   {"setmode",         set_mode},
diff --git a/src/ssl.lua b/src/ssl.lua
index 3d476c6..b182d53 100644
--- a/src/ssl.lua
+++ b/src/ssl.lua
@@ -203,16 +203,28 @@ local function newcontext(cfg)
 
    -- PSK
    if cfg.psk then
-      if type(cfg.psk) ~= "function" then
-         return nil, "invalid PSK callback parameter"
-      end
-
       if cfg.mode == "client" then
-         succ, msg = context.setclientpskcb(ctx, cfg.psk)
+         if type(cfg.psk) ~= "function" then
+            return nil, "invalid PSK configuration"
+         end
+         succ = context.setclientpskcb(ctx, cfg.psk)
          if not succ then return nil, msg end
       elseif cfg.mode == "server" then
-         succ, msg = context.setserverpskcb(ctx, cfg.psk)
-         if not succ then return nil, msg end
+         if type(cfg.psk) == "function" then
+            succ, msg = context.setserverpskcb(ctx, cfg.psk)
+            if not succ then return nil, msg end
+         elseif type(cfg.psk) == "table" then
+            if type(cfg.psk.hint) == "string" and type(cfg.psk.callback) == "function" then
+               succ, msg = context.setpskhint(ctx, cfg.psk.hint)
+               if not succ then return succ, msg end
+               succ = context.setserverpskcb(ctx, cfg.psk.callback)
+               if not succ then return succ, msg end
+            else
+               return nil, "invalid PSK configuration"
+            end
+         else
+            return nil, "invalid PSK configuration"
+         end
       end
    end