Skip to content

Commit 57c5d78

Browse files
DvdGiessenDvdGiessen
committed
extmod/modssl_mbedtls: Implement SSLSession support.
Signed-off-by: Daniël van de Giessen <daniel@dvdgiessen.nl>
1 parent 95d8b5f commit 57c5d78

File tree

1 file changed

+109
-5
lines changed

1 file changed

+109
-5
lines changed

‎extmod/modssl_mbedtls.c

Lines changed: 109 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,12 @@ typedef struct _mp_obj_ssl_context_t {
6161
int authmode;
6262
} mp_obj_ssl_context_t;
6363

64+
// This corresponds to an SSLSession object.
65+
typedef struct _mp_obj_ssl_session_t {
66+
mp_obj_base_t base;
67+
mbedtls_ssl_session session;
68+
} mp_obj_ssl_session_t;
69+
6470
// This corresponds to an SSLSocket object.
6571
typedef struct _mp_obj_ssl_socket_t {
6672
mp_obj_base_t base;
@@ -72,11 +78,12 @@ typedef struct _mp_obj_ssl_socket_t {
7278
int last_error; // The last error code, if any
7379
} mp_obj_ssl_socket_t;
7480

81+
STATIC const mp_obj_type_t ssl_session_type;
7582
STATIC const mp_obj_type_t ssl_context_type;
7683
STATIC const mp_obj_type_t ssl_socket_type;
7784

7885
STATIC mp_obj_t ssl_socket_make_new(mp_obj_ssl_context_t *ssl_context, mp_obj_t sock,
79-
bool server_side, bool do_handshake_on_connect, mp_obj_t server_hostname);
86+
bool server_side, bool do_handshake_on_connect, mp_obj_t server_hostname, mp_obj_t ssl_session);
8087

8188
/******************************************************************************/
8289
// Helper functions.
@@ -138,6 +145,60 @@ STATIC NORETURN void mbedtls_raise_error(int err) {
138145
#endif
139146
}
140147

148+
/******************************************************************************/
149+
// SSLSession type.
150+
151+
STATIC mp_obj_t ssl_session_make_new(const mp_obj_type_t *type_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
152+
mp_arg_check_num(n_args, n_kw, 1, 1, false);
153+
154+
mp_buffer_info_t bufinfo;
155+
mp_get_buffer_raise(args[0], &bufinfo, MP_BUFFER_READ);
156+
157+
mp_obj_ssl_session_t *self = m_new_obj(mp_obj_ssl_session_t);
158+
self->base.type = type_in;
159+
160+
mbedtls_ssl_session_init(&self->session);
161+
int ret = mbedtls_ssl_session_load(&self->session, bufinfo.buf, bufinfo.len);
162+
if (ret != 0) {
163+
mbedtls_raise_error(ret);
164+
}
165+
166+
return MP_OBJ_FROM_PTR(self);
167+
}
168+
169+
STATIC mp_obj_t ssl_session_serialize(mp_obj_t self_in) {
170+
mp_obj_ssl_session_t *self = MP_OBJ_TO_PTR(self_in);
171+
size_t len;
172+
vstr_t vstr;
173+
mbedtls_ssl_session_save(&self->session, NULL, 0, &len);
174+
vstr_init_len(&vstr, len);
175+
mbedtls_ssl_session_save(&self->session, (unsigned char*) vstr.buf, len, &len);
176+
return mp_obj_new_bytes_from_vstr(&vstr);
177+
}
178+
STATIC MP_DEFINE_CONST_FUN_OBJ_1(ssl_session_serialize_obj, ssl_session_serialize);
179+
180+
STATIC mp_int_t ssl_session_get_buffer(mp_obj_t self_in, mp_buffer_info_t *bufinfo, mp_uint_t flags) {
181+
if (flags != MP_BUFFER_READ) {
182+
return 1;
183+
}
184+
mp_get_buffer_raise(ssl_session_serialize(self_in), bufinfo, flags);
185+
return 0;
186+
}
187+
188+
STATIC const mp_rom_map_elem_t ssl_session_locals_dict_table[] = {
189+
{ MP_ROM_QSTR(MP_QSTR_serialize), MP_ROM_PTR(&ssl_session_serialize_obj) },
190+
};
191+
STATIC MP_DEFINE_CONST_DICT(ssl_session_locals_dict, ssl_session_locals_dict_table);
192+
193+
STATIC MP_DEFINE_CONST_OBJ_TYPE(
194+
ssl_session_type,
195+
MP_QSTR_SSLSession,
196+
MP_TYPE_FLAG_NONE,
197+
make_new, ssl_session_make_new,
198+
buffer, ssl_session_get_buffer,
199+
locals_dict, &ssl_session_locals_dict
200+
);
201+
141202
/******************************************************************************/
142203
// SSLContext type.
143204

@@ -277,11 +338,12 @@ STATIC void ssl_context_load_cadata(mp_obj_ssl_context_t *self, mp_obj_t cadata_
277338
}
278339

279340
STATIC mp_obj_t ssl_context_wrap_socket(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
280-
enum { ARG_server_side, ARG_do_handshake_on_connect, ARG_server_hostname };
341+
enum { ARG_server_side, ARG_do_handshake_on_connect, ARG_server_hostname, ARG_session };
281342
static const mp_arg_t allowed_args[] = {
282343
{ MP_QSTR_server_side, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
283344
{ MP_QSTR_do_handshake_on_connect, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = true} },
284345
{ MP_QSTR_server_hostname, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
346+
{ MP_QSTR_session, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
285347
};
286348

287349
// Parse arguments.
@@ -292,7 +354,7 @@ STATIC mp_obj_t ssl_context_wrap_socket(size_t n_args, const mp_obj_t *pos_args,
292354

293355
// Create and return the new SSLSocket object.
294356
return ssl_socket_make_new(self, sock, args[ARG_server_side].u_bool,
295-
args[ARG_do_handshake_on_connect].u_bool, args[ARG_server_hostname].u_obj);
357+
args[ARG_do_handshake_on_connect].u_bool, args[ARG_server_hostname].u_obj, args[ARG_session].u_obj);
296358
}
297359
STATIC MP_DEFINE_CONST_FUN_OBJ_KW(ssl_context_wrap_socket_obj, 2, ssl_context_wrap_socket);
298360

@@ -352,7 +414,7 @@ STATIC int _mbedtls_ssl_recv(void *ctx, byte *buf, size_t len) {
352414
}
353415

354416
STATIC mp_obj_t ssl_socket_make_new(mp_obj_ssl_context_t *ssl_context, mp_obj_t sock,
355-
bool server_side, bool do_handshake_on_connect, mp_obj_t server_hostname) {
417+
bool server_side, bool do_handshake_on_connect, mp_obj_t server_hostname, mp_obj_t ssl_session) {
356418

357419
// Verify the socket object has the full stream protocol
358420
mp_get_stream_raise(sock, MP_STREAM_OP_READ | MP_STREAM_OP_WRITE | MP_STREAM_OP_IOCTL);
@@ -384,6 +446,14 @@ STATIC mp_obj_t ssl_socket_make_new(mp_obj_ssl_context_t *ssl_context, mp_obj_t
384446
}
385447
}
386448

449+
if (ssl_session != mp_const_none) {
450+
mp_obj_ssl_session_t *session = MP_OBJ_TO_PTR(ssl_session);
451+
ret = mbedtls_ssl_set_session(&o->ssl, &session->session);
452+
if (ret != 0) {
453+
goto cleanup;
454+
}
455+
}
456+
387457
mbedtls_ssl_set_bio(&o->ssl, &o->sock, _mbedtls_ssl_send, _mbedtls_ssl_recv, NULL);
388458

389459
if (do_handshake_on_connect) {
@@ -553,6 +623,36 @@ STATIC mp_uint_t socket_ioctl(mp_obj_t o_in, mp_uint_t request, uintptr_t arg, i
553623
return ret;
554624
}
555625

626+
STATIC void ssl_socket_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
627+
mp_obj_ssl_socket_t *self = MP_OBJ_TO_PTR(self_in);
628+
if (dest[0] == MP_OBJ_NULL) {
629+
// Load attribute.
630+
if (attr == MP_QSTR_session) {
631+
mp_obj_ssl_session_t *o = m_new_obj(mp_obj_ssl_session_t);
632+
o->base.type = &ssl_session_type;
633+
mbedtls_ssl_session_init(&o->session);
634+
int ret = mbedtls_ssl_get_session(&self->ssl, &o->session);
635+
if (ret != 0) {
636+
mbedtls_raise_error(ret);
637+
}
638+
dest[0] = o;
639+
} else {
640+
// Continue lookup in locals_dict.
641+
dest[1] = MP_OBJ_SENTINEL;
642+
}
643+
} else if (dest[1] != MP_OBJ_NULL) {
644+
// Store attribute.
645+
if (attr == MP_QSTR_session) {
646+
mp_obj_ssl_session_t *ssl_session = MP_OBJ_TO_PTR(dest[1]);
647+
dest[0] = MP_OBJ_NULL;
648+
int ret = mbedtls_ssl_set_session(&self->ssl, &ssl_session->session);
649+
if (ret != 0) {
650+
mbedtls_raise_error(ret);
651+
}
652+
}
653+
}
654+
}
655+
556656
STATIC const mp_rom_map_elem_t ssl_socket_locals_dict_table[] = {
557657
{ MP_ROM_QSTR(MP_QSTR_read), MP_ROM_PTR(&mp_stream_read_obj) },
558658
{ MP_ROM_QSTR(MP_QSTR_readinto), MP_ROM_PTR(&mp_stream_readinto_obj) },
@@ -581,6 +681,7 @@ STATIC MP_DEFINE_CONST_OBJ_TYPE(
581681
MP_QSTR_SSLSocket,
582682
MP_TYPE_FLAG_NONE,
583683
protocol, &ssl_socket_stream_p,
684+
attr, ssl_socket_attr,
584685
locals_dict, &ssl_socket_locals_dict
585686
);
586687

@@ -596,6 +697,7 @@ STATIC mp_obj_t mod_ssl_wrap_socket(size_t n_args, const mp_obj_t *pos_args, mp_
596697
ARG_cert_reqs,
597698
ARG_cadata,
598699
ARG_do_handshake,
700+
ARG_session,
599701
};
600702
static const mp_arg_t allowed_args[] = {
601703
{ MP_QSTR_key, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
@@ -605,6 +707,7 @@ STATIC mp_obj_t mod_ssl_wrap_socket(size_t n_args, const mp_obj_t *pos_args, mp_
605707
{ MP_QSTR_cert_reqs, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = MBEDTLS_SSL_VERIFY_NONE}},
606708
{ MP_QSTR_cadata, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
607709
{ MP_QSTR_do_handshake, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = true} },
710+
{ MP_QSTR_session, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
608711
};
609712

610713
// Parse arguments.
@@ -633,7 +736,7 @@ STATIC mp_obj_t mod_ssl_wrap_socket(size_t n_args, const mp_obj_t *pos_args, mp_
633736

634737
// Create and return the new SSLSocket object.
635738
return ssl_socket_make_new(ssl_context, sock, args[ARG_server_side].u_bool,
636-
args[ARG_do_handshake].u_bool, args[ARG_server_hostname].u_obj);
739+
args[ARG_do_handshake].u_bool, args[ARG_server_hostname].u_obj, args[ARG_session].u_obj);
637740
}
638741
STATIC MP_DEFINE_CONST_FUN_OBJ_KW(mod_ssl_wrap_socket_obj, 1, mod_ssl_wrap_socket);
639742

@@ -645,6 +748,7 @@ STATIC const mp_rom_map_elem_t mp_module_ssl_globals_table[] = {
645748

646749
// Classes.
647750
{ MP_ROM_QSTR(MP_QSTR_SSLContext), MP_ROM_PTR(&ssl_context_type) },
751+
{ MP_ROM_QSTR(MP_QSTR_SSLSession), MP_ROM_PTR(&ssl_session_type) },
648752

649753
// Constants.
650754
{ MP_ROM_QSTR(MP_QSTR_PROTOCOL_TLS_CLIENT), MP_ROM_INT(MBEDTLS_SSL_IS_CLIENT) },

0 commit comments

Comments
 (0)