From fe4748f583389782b0a0fa5202612e12dc5fe22f Mon Sep 17 00:00:00 2001
From: Ian Kent <email@iankent.co.uk>
Date: Sun, 23 Nov 2014 04:21:53 +0000
Subject: [PATCH] Use io.ReadWriteCloser in place of net.TCPConn for SMTP
 session

---
 bindata.go                     |  2 +-
 mailhog/smtp/server/session.go | 27 ++++++++++++++-------------
 main.go                        |  3 ++-
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/bindata.go b/bindata.go
index b1e3966..b1021e1 100644
--- a/bindata.go
+++ b/bindata.go
@@ -1603,8 +1603,8 @@ var _bintree = &_bintree_t{nil, map[string]*_bintree_t{
 			"strutil.js":     &_bintree_t{assets_js_strutil_js, map[string]*_bintree_t{}},
 		}},
 		"templates": &_bintree_t{nil, map[string]*_bintree_t{
-			"layout.html": &_bintree_t{assets_templates_layout_html, map[string]*_bintree_t{}},
 			"index.html":  &_bintree_t{assets_templates_index_html, map[string]*_bintree_t{}},
+			"layout.html": &_bintree_t{assets_templates_layout_html, map[string]*_bintree_t{}},
 		}},
 	}},
 }}
diff --git a/mailhog/smtp/server/session.go b/mailhog/smtp/server/session.go
index da53589..51d263d 100644
--- a/mailhog/smtp/server/session.go
+++ b/mailhog/smtp/server/session.go
@@ -3,8 +3,8 @@ package server
 // http://www.rfc-editor.org/rfc/rfc5321.txt
 
 import (
+	"io"
 	"log"
-	"net"
 	"strings"
 
 	"github.com/ian-kent/Go-MailHog/mailhog/config"
@@ -14,18 +14,19 @@ import (
 
 // Session represents a SMTP session using net.TCPConn
 type Session struct {
-	conn  *net.TCPConn
-	proto *protocol.Protocol
-	conf  *config.Config
-	isTLS bool
-	line  string
+	conn          io.ReadWriteCloser
+	proto         *protocol.Protocol
+	conf          *config.Config
+	remoteAddress string
+	isTLS         bool
+	line          string
 }
 
-// Accept starts a new SMTP session using net.TCPConn
-func Accept(conn *net.TCPConn, conf *config.Config) {
+// Accept starts a new SMTP session using io.ReadWriteCloser
+func Accept(remoteAddress string, conn io.ReadWriteCloser, conf *config.Config) {
 	proto := protocol.NewProtocol()
 	proto.Hostname = conf.Hostname
-	session := &Session{conn, proto, conf, false, ""}
+	session := &Session{conn, proto, conf, remoteAddress, false, ""}
 	proto.LogHandler = session.logf
 	proto.MessageReceivedHandler = session.acceptMessage
 	proto.ValidateSenderHandler = session.validateSender
@@ -60,14 +61,14 @@ func (c *Session) acceptMessage(msg *data.Message) (id string, err error) {
 
 func (c *Session) logf(message string, args ...interface{}) {
 	message = strings.Join([]string{"[SMTP %s]", message}, " ")
-	args = append([]interface{}{c.conn.RemoteAddr()}, args...)
+	args = append([]interface{}{c.remoteAddress}, args...)
 	log.Printf(message, args...)
 }
 
 // Read reads from the underlying net.TCPConn
 func (c *Session) Read() bool {
 	buf := make([]byte, 1024)
-	n, err := c.conn.Read(buf)
+	n, err := io.Reader(c.conn).Read(buf)
 
 	if n == 0 {
 		c.logf("Connection closed by remote host\n")
@@ -91,7 +92,7 @@ func (c *Session) Read() bool {
 	if reply != nil {
 		c.Write(reply)
 		if reply.Status == 221 {
-			c.conn.Close()
+			io.Closer(c.conn).Close()
 		}
 	}
 
@@ -105,6 +106,6 @@ func (c *Session) Write(reply *protocol.Reply) {
 		logText := strings.Replace(l, "\n", "\\n", -1)
 		logText = strings.Replace(logText, "\r", "\\r", -1)
 		c.logf("Sent %d bytes: '%s'", len(l), logText)
-		c.conn.Write([]byte(l))
+		io.Writer(c.conn).Write([]byte(l))
 	}
 }
diff --git a/main.go b/main.go
index c42d772..e1431e3 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"flag"
+	"io"
 	"net"
 	"os"
 
@@ -122,6 +123,6 @@ func smtp_listen() *net.TCPListener {
 		}
 		defer conn.Close()
 
-		go smtp.Accept(conn.(*net.TCPConn), conf)
+		go smtp.Accept(conn.(*net.TCPConn).RemoteAddr().String(), io.ReadWriteCloser(conn), conf)
 	}
 }