2 files changed, 39 insertions, 18 deletions

diff --git a/src/internal/ptycon_driver_impl.h b/src/internal/ptycon_driver_impl.h
index a51c5a1..c46e146 100644
--- a/src/internal/ptycon_driver_impl.h
+++ b/src/internal/ptycon_driver_impl.h
@@ -31,4 +31,16 @@ struct ptyc_driver_ctx_impl {
 	struct ptyc_driver_ctx	ctx;
 };
 
+static inline struct ptyc_driver_ctx_impl * ptyc_get_driver_ictx(struct ptyc_driver_ctx * dctx)
+{
+	uintptr_t addr;
+
+	if (dctx) {
+		addr = (uintptr_t)dctx - offsetof(struct ptyc_driver_ctx_impl,ctx);
+		return (struct ptyc_driver_ctx_impl *)addr;
+	}
+
+	return 0;
+}
+
 #endif
diff --git a/src/pty/ptyc_pty_ctx.c b/src/pty/ptyc_pty_ctx.c
index 31a9c26..d796138 100644
--- a/src/pty/ptyc_pty_ctx.c
+++ b/src/pty/ptyc_pty_ctx.c
@@ -28,15 +28,16 @@ static int32_t ptyc_pty_open(nt_pty * hptm, nt_guid * pty_guid, nt_pty ** hpty)
 		&oa,0,0,0);
 }
 
-static int32_t ptyc_ptm_open(const struct ptyc_common_ctx * cctx)
+static int32_t ptyc_ptm_open(struct ptyc_common_ctx * cctx)
 {
 	nt_guid pty_guid = TTY_PTM_GUID;
+
 	return ptyc_pty_open(
 		0,&pty_guid,
-		(nt_pty **)&cctx->hptm);
+		&cctx->hptm);
 }
 
-static int32_t ptyc_pts_open(const struct ptyc_common_ctx * cctx)
+static int32_t ptyc_pts_open(struct ptyc_common_ctx * cctx)
 {
 	int32_t status;
 	nt_guid pty_guid = TTY_PTS_GUID;
@@ -46,27 +47,34 @@ static int32_t ptyc_pts_open(const struct ptyc_common_ctx * cctx)
 
 	return ptyc_pty_open(
 		cctx->hptm,&pty_guid,
-		(nt_pty **)&cctx->hpts);
+		&cctx->hpts);
 }
 
 int ptyc_alloc_pty(struct ptyc_driver_ctx * dctx)
 {
-	int32_t status;
+	int32_t				status;
+	struct ptyc_driver_ctx_impl *	ictx;
+	struct ptyc_common_ctx *	cctx;
+
+	if (!(ictx = ptyc_get_driver_ictx(dctx)))
+		return NT_STATUS_INVALID_HANDLE;
+
+	cctx = &ictx->cctx;
 
-	if (dctx->cctx->hpts || dctx->cctx->hptm)
+	if (cctx->hpts || cctx->hptm)
 		return ptyc_set_status(
 			dctx,NT_STATUS_DEVICE_ALREADY_ATTACHED);
 
-	if ((status = ptyc_ptm_open(dctx->cctx)))
+	if ((status = ptyc_ptm_open(cctx)))
 		return ptyc_set_status(
 			dctx,status);
 
-	if (!(dctx->cctx->drvflags & PTYC_DRIVER_DBG_OVEN)
-			&& !(dctx->cctx->drvflags & PTYC_DRIVER_DBG_RAW))
+	if (!(cctx->drvflags & PTYC_DRIVER_DBG_OVEN)
+			&& !(cctx->drvflags & PTYC_DRIVER_DBG_RAW))
 		return ptyc_set_status(
 			dctx,NT_STATUS_SUCCESS);
 
-	if ((status = ptyc_pts_open(dctx->cctx)))
+	if ((status = ptyc_pts_open(cctx)))
 		return ptyc_set_status(
 			dctx,status);
 
@@ -76,17 +84,18 @@ int ptyc_alloc_pty(struct ptyc_driver_ctx * dctx)
 
 void ptyc_free_pty(struct ptyc_driver_ctx * dctx)
 {
-	struct ptyc_common_ctx * cctx;
+	struct ptyc_driver_ctx_impl * ictx;
 
-	cctx = (struct ptyc_common_ctx *)dctx->cctx;
+	if (!(ictx = ptyc_get_driver_ictx(dctx)))
+		return;
 
-	if (cctx->hpts) {
-		ntapi->pty_close(cctx->hpts);
-		cctx->hpts = 0;
+	if (ictx->cctx.hpts) {
+		ntapi->pty_close(ictx->cctx.hpts);
+		ictx->cctx.hpts = 0;
 	}
 
-	if (cctx->hptm) {
-		ntapi->pty_close(cctx->hptm);
-		cctx->hptm = 0;
+	if (ictx->cctx.hptm) {
+		ntapi->pty_close(ictx->cctx.hptm);
+		ictx->cctx.hptm = 0;
 	}
 }