From 9e9815ae5a6771d2002024b3bae8767b295b6319 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Thu, 11 Apr 2013 21:13:47 -0500
Subject: [PATCH] Added plain text password authentication

---
 conn.go        | 24 +++++++++++++++++++++++-
 conn_test.go   | 21 +++++++++++++++++----
 test_setup.sql |  6 ++++--
 3 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/conn.go b/conn.go
index d0fcbea4..91abbd83 100644
--- a/conn.go
+++ b/conn.go
@@ -14,6 +14,7 @@ type conn struct {
 	pid           int32             // backend pid
 	secretKey     int32             // key to use to send a cancel query message to the server
 	runtimeParams map[string]string // parameters that have been reported by the server
+	options       map[string]string // options used when establishing connection
 	txStatus      byte
 }
 
@@ -23,6 +24,11 @@ type conn struct {
 func Connect(options map[string]string) (c *conn, err error) {
 	c = new(conn)
 
+	c.options = make(map[string]string)
+	for k, v := range options {
+		c.options[k] = v
+	}
+
 	var present bool
 	var socket string
 
@@ -55,7 +61,9 @@ func Connect(options map[string]string) (c *conn, err error) {
 			case backendKeyData:
 				c.rxBackendKeyData(r)
 			case authenticationX:
-				c.rxAuthenticationX(r)
+				if err = c.rxAuthenticationX(r); err != nil {
+					return nil, err
+				}
 			case readyForQuery:
 				return c, nil
 			default:
@@ -179,6 +187,8 @@ func (c *conn) rxAuthenticationX(r *messageReader) (err error) {
 	code := r.readInt32()
 	switch code {
 	case 0: // AuthenticationOk
+	case 3: // AuthenticationCleartextPassword
+		c.txPasswordMessage(c.options["password"])
 	default:
 		err = errors.New("Received unknown authentication message")
 	}
@@ -257,6 +267,18 @@ func (c *conn) txStartupMessage(msg *startupMessage) (err error) {
 	return
 }
 
+func (c *conn) txPasswordMessage(password string) (err error) {
+	bufSize := 5 + len(password) + 1 // message identifier (1), message size (4), password, null string terminator (1)
+	buf := c.getBuf(bufSize)
+	buf[0] = 'p'
+	binary.BigEndian.PutUint32(buf[1:5], uint32(bufSize-1))
+	copy(buf[5:], password)
+	buf[bufSize-1] = 0
+
+	_, err = c.conn.Write(buf)
+	return err
+}
+
 // Gets a []byte of n length. If possible it will reuse the connection buffer
 // otherwise it will allocate a new buffer
 func (c *conn) getBuf(n int) (buf []byte) {
diff --git a/conn_test.go b/conn_test.go
index 5c7a6f2d..bbbcc710 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -5,7 +5,7 @@ import (
 )
 
 func TestConnect(t *testing.T) {
-	conn, err := Connect(map[string]string{"socket": "/private/tmp/.s.PGSQL.5432", "user": "pgx", "database": "pgx_test"})
+	conn, err := Connect(map[string]string{"socket": "/private/tmp/.s.PGSQL.5432", "user": "pgx_none", "database": "pgx_test"})
 	if err != nil {
 		t.Fatal("Unable to establish connection")
 	}
@@ -29,8 +29,8 @@ func TestConnect(t *testing.T) {
 	}
 
 	rows, err = conn.Query("select current_user")
-	if err != nil || rows[0]["current_user"] != "pgx" {
-		t.Error("Did not connect as specified user (pgx)")
+	if err != nil || rows[0]["current_user"] != "pgx_none" {
+		t.Error("Did not connect as specified user (pgx_none)")
 	}
 
 	err = conn.Close()
@@ -46,8 +46,21 @@ func TestConnectWithInvalidUser(t *testing.T) {
 		t.Fatal("Did not receive expected error when connecting with invalid user")
 	}
 }
+
+func TestConnectWithPassword(t *testing.T) {
+	conn, err := Connect(map[string]string{"socket": "/private/tmp/.s.PGSQL.5432", "user": "pgx_pw", "password": "secret", "database": "pgx_test"})
+	if err != nil {
+		t.Fatal("Unable to establish connection: " + err.Error())
+	}
+
+	err = conn.Close()
+	if err != nil {
+		t.Fatal("Unable to close connection")
+	}
+}
+
 func TestQuery(t *testing.T) {
-	conn, err := Connect(map[string]string{"socket": "/private/tmp/.s.PGSQL.5432", "user": "pgx", "database": "pgx_test"})
+	conn, err := Connect(map[string]string{"socket": "/private/tmp/.s.PGSQL.5432", "user": "pgx_none", "database": "pgx_test"})
 	if err != nil {
 		t.Fatal("Unable to establish connection")
 	}
diff --git a/test_setup.sql b/test_setup.sql
index 18bc364c..15af361c 100644
--- a/test_setup.sql
+++ b/test_setup.sql
@@ -1,5 +1,7 @@
 drop database if exists pgx_test;
-drop user if exists pgx;
+drop user if exists pgx_none;
+drop user if exists pgx_pw;
 
-create user pgx;
+create user pgx_none;
+create user pgx_pw password 'secret';
 create database pgx_test;