Refactor the TLS code so that the push/pull functions are aware of their

containing stream object. This is in preparation for supporting GnuTLS 3.x.


git-svn-id: svn://svn.code.sf.net/p/tigervnc/code/trunk@4921 3789f03b-4d11-0410-bbf8-ca57d06f2519

6 files changed, 43 insertions, 29 deletions

diff --git a/common/rdr/TLSInStream.cxx b/common/rdr/TLSInStream.cxx
index ddc99917..e553085f 100644
--- a/common/rdr/TLSInStream.cxx
+++ b/common/rdr/TLSInStream.cxx
@@ -36,10 +36,10 @@ using namespace rdr;
 
 enum { DEFAULT_BUF_SIZE = 16384 };
 
-ssize_t rdr::gnutls_InStream_pull(gnutls_transport_ptr str, void* data,
-				  size_t size)
+ssize_t TLSInStream::pull(gnutls_transport_ptr str, void* data, size_t size)
 {
-  InStream* in= (InStream*) str;
+  TLSInStream* self= (TLSInStream*) str;
+  InStream *in = self->in;
 
   try {
     if (!in->check(1, 1, false)) {
@@ -63,11 +63,19 @@ ssize_t rdr::gnutls_InStream_pull(gnutls_transport_ptr str, void* data,
 TLSInStream::TLSInStream(InStream* _in, gnutls_session _session)
   : session(_session), in(_in), bufSize(DEFAULT_BUF_SIZE), offset(0)
 {
+  gnutls_transport_ptr recv, send;
+
   ptr = end = start = new U8[bufSize];
+
+  gnutls_transport_set_pull_function(session, pull);
+  gnutls_transport_get_ptr2(session, &recv, &send);
+  gnutls_transport_set_ptr2(session, this, send);
 }
 
 TLSInStream::~TLSInStream()
 {
+  gnutls_transport_set_pull_function(session, NULL);
+
   delete[] start;
 }
 
diff --git a/common/rdr/TLSInStream.h b/common/rdr/TLSInStream.h
index 7fad6fa5..65a783c8 100644
--- a/common/rdr/TLSInStream.h
+++ b/common/rdr/TLSInStream.h
@@ -41,6 +41,7 @@ namespace rdr {
   private:
     int overrun(int itemSize, int nItems, bool wait);
     int readTLS(U8* buf, int len, bool wait);
+    static ssize_t pull(gnutls_transport_ptr str, void* data, size_t size);
 
     gnutls_session session;
     InStream* in;
@@ -48,9 +49,6 @@ namespace rdr {
     int offset;
     U8* start;
   };
-
-  ssize_t gnutls_InStream_pull(gnutls_transport_ptr,void*, size_t);
-
 };
 
 #endif
diff --git a/common/rdr/TLSOutStream.cxx b/common/rdr/TLSOutStream.cxx
index d577ccc9..ec21670c 100644
--- a/common/rdr/TLSOutStream.cxx
+++ b/common/rdr/TLSOutStream.cxx
@@ -36,10 +36,11 @@ using namespace rdr;
 
 enum { DEFAULT_BUF_SIZE = 16384 };
 
-ssize_t rdr::gnutls_OutStream_push(gnutls_transport_ptr str, const void* data,
+ssize_t TLSOutStream::push(gnutls_transport_ptr str, const void* data,
 				   size_t size)
 {
-  OutStream* out = (OutStream*) str;
+  TLSOutStream* self= (TLSOutStream*) str;
+  OutStream *out = self->out;
 
   try {
     out->writeBytes(data, size);
@@ -55,8 +56,14 @@ ssize_t rdr::gnutls_OutStream_push(gnutls_transport_ptr str, const void* data,
 TLSOutStream::TLSOutStream(OutStream* _out, gnutls_session _session)
   : session(_session), out(_out), bufSize(DEFAULT_BUF_SIZE), offset(0)
 {
+  gnutls_transport_ptr recv, send;
+
   ptr = start = new U8[bufSize];
   end = start + bufSize;
+
+  gnutls_transport_set_push_function(session, push);
+  gnutls_transport_get_ptr2(session, &recv, &send);
+  gnutls_transport_set_ptr2(session, recv, this);
 }
 
 TLSOutStream::~TLSOutStream()
@@ -67,6 +74,8 @@ TLSOutStream::~TLSOutStream()
   } catch (Exception&) {
   }
 #endif
+  gnutls_transport_set_push_function(session, NULL);
+
   delete [] start;
 }
 
diff --git a/common/rdr/TLSOutStream.h b/common/rdr/TLSOutStream.h
index 5eb512ec..a291f426 100644
--- a/common/rdr/TLSOutStream.h
+++ b/common/rdr/TLSOutStream.h
@@ -43,6 +43,7 @@ namespace rdr {
 
   private:
     int writeTLS(const U8* data, int length);
+    static ssize_t push(gnutls_transport_ptr str, const void* data, size_t size);
 
     gnutls_session session;
     OutStream* out;
@@ -50,8 +51,6 @@ namespace rdr {
     U8* start;
     int offset;
   };
-
-  ssize_t gnutls_OutStream_push(gnutls_transport_ptr, const void*, size_t);
 };
 
 #endif
diff --git a/common/rfb/CSecurityTLS.cxx b/common/rfb/CSecurityTLS.cxx
index 6b6b017c..59fd5067 100644
--- a/common/rfb/CSecurityTLS.cxx
+++ b/common/rfb/CSecurityTLS.cxx
@@ -188,20 +188,20 @@ bool CSecurityTLS::processMsg(CConnection* cc)
       throw AuthFailureException("gnutls_set_default_priority failed");
 
     setParam();
-    
-    gnutls_transport_set_pull_function(session, rdr::gnutls_InStream_pull);
-    gnutls_transport_set_push_function(session, rdr::gnutls_OutStream_push);
-    gnutls_transport_set_ptr2(session,
-			      (gnutls_transport_ptr) is,
-			      (gnutls_transport_ptr) os);
   }
 
+  rdr::TLSInStream *tlsis = new rdr::TLSInStream(is, session);
+  rdr::TLSOutStream *tlsos = new rdr::TLSOutStream(os, session);
+
   int err;
   err = gnutls_handshake(session);
-  if (err != GNUTLS_E_SUCCESS && !gnutls_error_is_fatal(err))
-    return false;
-
   if (err != GNUTLS_E_SUCCESS) {
+    delete tlsis;
+    delete tlsos;
+
+    if (!gnutls_error_is_fatal(err))
+      return false;
+
     vlog.error("TLS Handshake failed: %s\n", gnutls_strerror (err));
     shutdown(false);
     throw AuthFailureException("TLS Handshake failed");
@@ -209,8 +209,7 @@ bool CSecurityTLS::processMsg(CConnection* cc)
 
   checkSession();
 
-  cc->setStreams(fis = new rdr::TLSInStream(is, session),
-		 fos = new rdr::TLSOutStream(os, session));
+  cc->setStreams(fis = tlsis, fos = tlsos);
 
   return true;
 }
diff --git a/common/rfb/SSecurityTLS.cxx b/common/rfb/SSecurityTLS.cxx
index 9a34f3a7..7a1a41ef 100644
--- a/common/rfb/SSecurityTLS.cxx
+++ b/common/rfb/SSecurityTLS.cxx
@@ -148,17 +148,19 @@ bool SSecurityTLS::processMsg(SConnection *sc)
       throw;
     }
 
-    gnutls_transport_set_pull_function(session,rdr::gnutls_InStream_pull);
-    gnutls_transport_set_push_function(session,rdr::gnutls_OutStream_push);
-    gnutls_transport_set_ptr2(session,
-			      (gnutls_transport_ptr)is,
-			      (gnutls_transport_ptr)os);
     os->writeU8(1);
     os->flush();
   }
 
+  rdr::TLSInStream *tlsis = new rdr::TLSInStream(is, session);
+  rdr::TLSOutStream *tlsos = new rdr::TLSOutStream(os, session);
+
   int err;
-  if ((err = gnutls_handshake(session)) != GNUTLS_E_SUCCESS) {
+  err = gnutls_handshake(session);
+  if (err != GNUTLS_E_SUCCESS) {
+    delete tlsis;
+    delete tlsos;
+
     if (!gnutls_error_is_fatal(err)) {
       vlog.debug("Deferring completion of TLS handshake: %s", gnutls_strerror(err));
       return false;
@@ -170,8 +172,7 @@ bool SSecurityTLS::processMsg(SConnection *sc)
 
   vlog.debug("Handshake completed");
 
-  sc->setStreams(fis=new rdr::TLSInStream(is,session),
-		 fos=new rdr::TLSOutStream(os,session));
+  sc->setStreams(fis = tlsis, fos = tlsos);
 
   return true;
 }