Reject invalid status codes

5 files changed, 21 insertions, 12 deletions

diff --git a/client.go b/client.go
index 2914d41..c067dc2 100644
--- a/client.go
+++ b/client.go
@@ -15,11 +15,11 @@ import (
 
 // Client errors.
 var (
-	ErrProtocol              = errors.New("gemini: protocol error")
-	ErrInvalidURL            = errors.New("gemini: requested URL is invalid")
-	ErrCertificateNotValid   = errors.New("gemini: certificate is invalid")
+	ErrInvalidURL            = errors.New("gemini: invalid URL")
+	ErrInvalidResponse       = errors.New("gemini: invalid response")
+	ErrInvalidCertificate    = errors.New("gemini: invalid certificate")
+	ErrUnknownCertificate    = errors.New("gemini: unknown certificate")
 	ErrCertificateNotTrusted = errors.New("gemini: certificate is not trusted")
-	ErrCertificateUnknown    = errors.New("gemini: certificate is unknown")
 )
 
 // Request represents a Gemini request.
@@ -133,11 +133,18 @@ func (resp *Response) read(r *bufio.Reader) error {
 	}
 	resp.Status = status
 
+	// Disregard invalid status codes
+	const minStatus, maxStatus = 1, 6
+	statusClass := status / 10
+	if statusClass < minStatus || statusClass > maxStatus {
+		return ErrInvalidResponse
+	}
+
 	// Read one space
 	if b, err := r.ReadByte(); err != nil {
 		return err
 	} else if b != ' ' {
-		return ErrProtocol
+		return ErrInvalidResponse
 	}
 
 	// Read the meta
@@ -149,7 +156,7 @@ func (resp *Response) read(r *bufio.Reader) error {
 	meta = meta[:len(meta)-1]
 	// Ensure meta is less than or equal to 1024 bytes
 	if len(meta) > 1024 {
-		return ErrProtocol
+		return ErrInvalidResponse
 	}
 	resp.Meta = meta
 
@@ -157,7 +164,7 @@ func (resp *Response) read(r *bufio.Reader) error {
 	if b, err := r.ReadByte(); err != nil {
 		return err
 	} else if b != '\n' {
-		return ErrProtocol
+		return ErrInvalidResponse
 	}
 
 	// Read response body
@@ -244,11 +251,11 @@ func (c *Client) Send(req *Request) (*Response, error) {
 	// Read the response
 	resp := &Response{}
 	r := bufio.NewReader(conn)
-	// Store connection information
-	resp.TLS = conn.ConnectionState()
 	if err := resp.read(r); err != nil {
 		return nil, err
 	}
+	// Store connection information
+	resp.TLS = conn.ConnectionState()
 	return resp, nil
 }
 
diff --git a/examples/auth/auth.go b/examples/auth/auth.go
index f135c48..0c99535 100644
--- a/examples/auth/auth.go
+++ b/examples/auth/auth.go
@@ -38,6 +38,7 @@ func main() {
 	// To generate a TLS key pair, run:
 	//
 	//     go run -tags=example ../cert
+	//
 	cert, err := tls.LoadX509KeyPair("examples/client/localhost.crt", "examples/client/localhost.key")
 	if err != nil {
 		log.Fatal(err)
diff --git a/examples/client/client.go b/examples/client/client.go
index 402e8f6..30a0a9b 100644
--- a/examples/client/client.go
+++ b/examples/client/client.go
@@ -28,7 +28,7 @@ func init() {
 				// Alert the user that the certificate is not trusted
 				fmt.Printf("Warning: Certificate for %s is not trusted!\n", hostname)
 				fmt.Println("This could indicate a Man-in-the-Middle attack.")
-			case gemini.ErrCertificateUnknown:
+			case gemini.ErrUnknownCertificate:
 				// Prompt the user to trust the certificate
 				trust := trustCertificate(cert)
 				switch trust {
diff --git a/examples/server/server.go b/examples/server/server.go
index bda87e6..e761cd0 100644
--- a/examples/server/server.go
+++ b/examples/server/server.go
@@ -14,6 +14,7 @@ func main() {
 	// To generate a TLS key pair, run:
 	//
 	//     go run -tags=example ../cert
+	//
 	cert, err := tls.LoadX509KeyPair("examples/server/localhost.crt", "examples/server/localhost.key")
 	if err != nil {
 		log.Fatal(err)
diff --git a/tofu.go b/tofu.go
index 7da58dc..21384ab 100644
--- a/tofu.go
+++ b/tofu.go
@@ -79,7 +79,7 @@ func (k *KnownHosts) AddTemporary(hostname string, cert *x509.Certificate) {
 // Lookup looks for the provided certificate in the list of known hosts.
 // If the hostname is in the list, but the fingerprint differs,
 // Lookup returns ErrCertificateNotTrusted.
-// If the hostname is not in the list, Lookup returns ErrCertificateUnknown.
+// If the hostname is not in the list, Lookup returns ErrUnknownCertificate.
 // If the certificate is found and the fingerprint matches, error will be nil.
 func (k *KnownHosts) Lookup(hostname string, cert *x509.Certificate) error {
 	now := time.Now().Unix()
@@ -99,7 +99,7 @@ func (k *KnownHosts) Lookup(hostname string, cert *x509.Certificate) error {
 		// Fingerprint does not match
 		return ErrCertificateNotTrusted
 	}
-	return ErrCertificateUnknown
+	return ErrUnknownCertificate
 }
 
 // Parse parses the provided reader and adds the parsed known hosts to the list.