3 files changed, 35 insertions, 50 deletions

diff --git a/gemini/gemini.go b/gemini/gemini.go
index 8f4c6ef..bd341e4 100644
--- a/gemini/gemini.go
+++ b/gemini/gemini.go
@@ -12,53 +12,47 @@ import (
 const GeminiNetwork = "tcp"
 const GeminiPort = 1965
 const GeminiPrefix = "gemini://"
+const requestSuffix = "\r\n"
 
-func FormatRequest(request string) []byte {
-	if !strings.HasPrefix(request, GeminiPrefix) {
-		request = GeminiPrefix + request
-	}
-	request = request + "\r\n"
-	return []byte(request)
-}
-
-type GeminiConn struct {
-	host   string
-	config *tls.Config
+type GeminiURI struct {
+	host    string
+	request []byte
 }
 
-func NewGeminiConn(host string, port int) *GeminiConn {
-	conf := tls.Config{
-		MinVersion:         tls.VersionTLS12,
-		InsecureSkipVerify: true,
+func Parse(request, host string, port int) (*GeminiURI, error) {
+	// add gemini:// scheme if missing
+	if !strings.HasPrefix(request, GeminiPrefix) {
+		request = GeminiPrefix + request
 	}
 
+	// create GeminiURI if host has been specified
 	if strings.Contains(host, ":") {
-		return &GeminiConn{host, &conf}
-	} else {
-		return &GeminiConn{host + ":" + strconv.Itoa(port), &conf}
-	}
-}
-
-func NewGeminiConnFromRequest(request string, port int) (*GeminiConn, error) {
-	if !strings.HasPrefix(request, GeminiPrefix) {
-		request = GeminiPrefix + request
+		return &GeminiURI{host, []byte(request + requestSuffix)}, nil
+	} else if host != "" {
+		return &GeminiURI{host + ":" + strconv.Itoa(port), []byte(request + requestSuffix)}, nil
 	}
 
+	// get host from request
 	u, err := url.Parse(request)
 	if err != nil {
 		return nil, err
 	}
-	return NewGeminiConn(u.Host, port), nil
+
+	if strings.Contains(u.Host, ":") {
+		return &GeminiURI{u.Host, []byte(request + requestSuffix)}, nil
+	} else {
+		return &GeminiURI{u.Host + ":" + strconv.Itoa(port), []byte(request + requestSuffix)}, nil
+	}
 }
 
-func (c *GeminiConn) Get(request []byte, out *os.File) error {
-	conn, err := tls.Dial(GeminiNetwork, c.host, c.config)
+func (c *GeminiURI) Get(config *tls.Config, out *os.File) error {
+	conn, err := tls.Dial(GeminiNetwork, c.host, config)
 	if err != nil {
 		return err
 	}
 	defer conn.Close()
 
-	n, err := conn.Write(request)
+	n, err := conn.Write(c.request)
 	if err != nil {
 		return err
 	}
diff --git a/gemini/gemini_test.go b/gemini/gemini_test.go
index ffd2588..1d69498 100644
--- a/gemini/gemini_test.go
+++ b/gemini/gemini_test.go
@@ -1,43 +1,28 @@
 package gemini
 
 import (
-	"bytes"
 	"strconv"
 	"testing"
 )
 
-func TestNewGeminiConn(t *testing.T) {
+func TestParse(t *testing.T) {
 	tables := []struct {
-		url  string
-		port int
 		host string
+		port int
+		url string
 	}{
 		{"hostname.com", GeminiPort, "hostname.com:" + strconv.Itoa(GeminiPort)},
 		{"hostname.com", 1234, "hostname.com:1234"},
 	}
 
 	for _, table := range tables {
-		conn, err := NewGeminiConnFromRequest(table.url, table.port)
+		conn, err := Parse("", table.host, table.port)
 		if err != nil {
 			t.Fatalf("NewGeminiConn error: %s", err.Error())
 		}
-		if conn.host != table.host {
+		if conn.host != table.url {
 			t.Fatalf("NewGeminiConn error: wrong hostname %s", conn.host)
 		}
 	}
 }
 
-func TestFormatRequest(t *testing.T) {
-	tables := []struct {
-		input  string
-		output []byte
-	}{
-		{"hostname.com", []byte("gemini://hostname.com\r\n")},
-	}
-
-	for _, table := range tables {
-		if !bytes.Equal(FormatRequest(table.input), table.output) {
-			t.Fatalf("FormatRequest failed on: %s\n", table.input)
-		}
-	}
-}
diff --git a/main.go b/main.go
index 3cfacac..f008fa3 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"crypto/tls"
 	"flag"
 	"fmt"
 	"gemini"
@@ -76,12 +77,17 @@ func main() {
 		log.Fatal(err)
 	}
 
+	conf := tls.Config{
+		MinVersion:         tls.VersionTLS12,
+		InsecureSkipVerify: true,
+	}
+
 	log.Printf("Got %d requests\n", len(requests))
 	for _, request := range requests {
 		log.Println(request)
-		if conn, err := gemini.NewGeminiConnFromRequest(request, flags.port); err != nil {
+		if conn, err := gemini.Parse(request, flags.host, flags.port); err != nil {
 			log.Print(err)
-		} else if err := conn.Get(gemini.FormatRequest(request), output); err != nil {
+		} else if err := conn.Get(&conf, output); err != nil {
 			log.Print(err)
 		}
 	}