refactor some stuff out of Client into Connection

so it can be shared with Server at some point

also do some error handling on startup message

pghottub/client.cpp

@@ -6,28 +6,20 @@
 
 #include <map>
 
-#include <mordor/assert.h>
 #include <mordor/endian.h>
 #include <mordor/log.h>
-#include <mordor/streams/buffered.h>
-#include <mordor/streams/limited.h>
+#include <mordor/streams/buffer.h>
 #include <mordor/streams/stream.h>
 
-#include "postgres.h"
-
 using namespace Mordor;
 
 static Logger::ptr g_log = Log::lookup("pghottub:client");
 
 namespace PgHotTub {
 
 Client::Client(Stream::ptr stream)
-{
-    MORDOR_ASSERT(stream->supportsRead());
-    MORDOR_ASSERT(stream->supportsWrite());
-    MORDOR_ASSERT(!stream->supportsSeek());
-    m_stream.reset(new BufferedStream(stream));
-}
+    : Connection(stream)
+{}
 
 void
 Client::run()
@@ -42,38 +34,37 @@ Client::run()
 void
 Client::startup()
 {
-    unsigned int length;
-    m_stream->read(&length, 4u);
-    length = byteswap(length);
-    // TODO: disconnect unless length >= 8
+    V2MessageType type;
+    Buffer message;
+
+    readV2Message(type, message);
 
-    unsigned int protocolVersion;
-    m_stream->read(&protocolVersion, 4u);
-    protocolVersion = byteswap(protocolVersion);
+    if (type == SSL_REQUEST) {
+        if (message.readAvailable() != 0)
+            MORDOR_THROW_EXCEPTION(InvalidProtocolException());
 
-    if (protocolVersion == SSL_REQUEST) {
         // TODO: wrap in SSL
         m_stream->write("N", 1u);
         m_stream->flush();
-        startup();
-        return;
-    }
-    if (protocolVersion != 0x00030000) {
-        // TODO: disconnect unless version 3
+        readV2Message(type, message);
     }
 
+    if (type != STARTUP_REQUEST_V3)
+        MORDOR_THROW_EXCEPTION(InvalidProtocolException());
+
     std::map<std::string, std::string> parameters;
     while (true) {
-        std::string name = m_stream->getDelimited('\0');
+        std::string name = message.getDelimited('\0');
         if (name.length() == 1)
             break;
         name.resize(name.size() - 1);
-        std::string value = m_stream->getDelimited('\0', false, false);
+        std::string value = message.getDelimited('\0', false, false);
 
         parameters[name] = value;
     }
 
-    // TODO: check stream is at EOF
+    if (message.readAvailable() != 0)
+        MORDOR_THROW_EXCEPTION(InvalidProtocolException());
 }
 
 }

pghottub/client.h

@@ -2,16 +2,19 @@
 
 #include <boost/shared_ptr.hpp>
 
+#include "pghottub/connection.h"
+
 namespace Mordor {
 class Stream;
 }
 
 namespace PgHotTub {
 
-class Client
+class Client : public Connection
 {
 public:
     typedef boost::shared_ptr<Client> ptr;
+    struct InvalidProtocolException : virtual Mordor::Exception {};
 
 public:
     Client(boost::shared_ptr<Mordor::Stream> stream);
@@ -20,9 +23,6 @@ class Client
 
 private:
     void startup();
-
-private:
-    boost::shared_ptr<Mordor::Stream> m_stream;
 };
 
 }

pghottub/connection.cpp

@@ -0,0 +1,39 @@
+// Copyright (c) 2013 - Cody Cutrer
+
+#include <mordor/predef.h>
+
+#include "pghottub/connection.h"
+
+#include <mordor/assert.h>
+#include <mordor/streams/buffered.h>
+#include <mordor/endian.h>
+
+using namespace Mordor;
+
+namespace PgHotTub {
+
+Connection::Connection(Stream::ptr stream)
+{
+    MORDOR_ASSERT(stream->supportsRead());
+    MORDOR_ASSERT(stream->supportsWrite());
+    MORDOR_ASSERT(!stream->supportsSeek());
+    m_stream.reset(new BufferedStream(stream));
+}
+
+void
+Connection::readV2Message(V2MessageType &type, Buffer &message)
+{
+    unsigned int length;
+    m_stream->read(&length, 4u);
+    length = byteswap(length);
+    if (length < 8)
+        MORDOR_THROW_EXCEPTION(MessageTooShort());
+
+    m_stream->read(&type, 4u);
+    type = byteswap(type);
+
+    if (length > 8)
+        m_stream->read(message, length - 8);
+}
+
+}

pghottub/connection.h

@@ -0,0 +1,41 @@
+// Copyright (c) 2013 - Cody Cutrer
+
+#include <boost/noncopyable.hpp>
+
+#include <mordor/exception.h>
+
+namespace Mordor {
+struct Buffer;
+class Stream;
+}
+
+namespace PgHotTub {
+
+class Connection : boost::noncopyable
+{
+public:
+    struct MessageTooShort : virtual Mordor::Exception {};
+
+protected:
+    enum V2MessageType
+    {
+        SSL_REQUEST         = 80877103,
+        STARTUP_REQUEST_V2  = 0x00020000,
+        STARTUP_REQUEST_V3  = 0x00030000
+    };
+
+    enum V3MessageType
+    {
+    };
+
+protected:
+    Connection(boost::shared_ptr<Mordor::Stream> stream);
+
+    void readV2Message(V2MessageType &type, Mordor::Buffer &message);
+    void readV3Message(V3MessageType &type, Mordor::Buffer &message);
+
+protected:
+    boost::shared_ptr<Mordor::Stream> m_stream;
+};
+
+}

pghottub/pghottub.vcxproj

@@ -89,6 +89,7 @@
   </ItemDefinitionGroup>
   <ItemGroup>
     <ClCompile Include="client.cpp" />
+    <ClCompile Include="connection.cpp" />
     <ClCompile Include="hot_tub.cpp" />
     <ClCompile Include="listener.cpp" />
     <ClCompile Include="main.cpp" />
@@ -102,9 +103,9 @@
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="client.h" />
+    <ClInclude Include="connection.h" />
     <ClInclude Include="hot_tub.h" />
     <ClInclude Include="listener.h" />
-    <ClInclude Include="postgres.h" />
     <ClInclude Include="tcp_listener.h" />
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />