vendor/github.com/apigee-labs/transicator/pgclient/mockserver.go - apidVerifyApiKey - Git at Google

 /*
 Copyright 2017 The Transicator Authors

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
 package pgclient

 import (
 	"bytes"
 	"crypto/rand"
 	"crypto/tls"
 	"encoding/binary"
 	"fmt"
 	"io"
 	"net"
 	"regexp"

 	log "github.com/Sirupsen/logrus"
 )

 const (
 	mockUserName     = "mock"
 	mockPassword     = "mocketty"
 	mockDatabaseName = "turtle"
 )

 type mockState int

 const (
 	mockIdle mockState = 1
 )

 // MockAuth specifies what type of authentication the server supports
 type MockAuth int

 // Different auth types
 const (
 	MockTrust MockAuth = 0
 	MockClear MockAuth = 1
 	MockMD5   MockAuth = 2
 )

 var insertRE = regexp.MustCompile("insert into mock values \\('([\\w]+)', '([\\w]+)'\\)")

 /*
 A MockServer is a server that implements a little bit of the Postgres wire
 protocol. We can use it for testing of the wire protocol client. In particular,
 we can use it to test the myriad of password authentication and TLS options
 without having to start and stop a real Postgres server in the test suite.
 */
 type MockServer struct {
 	listener  net.Listener
 	mockTable map[string]string
 	authType  MockAuth
 	tlsConfig *tls.Config
 	forceTLS  bool
 }

 /*
 NewMockServer starts a new server in the current process, listening on the
 specified port.
 */
 func NewMockServer() *MockServer {
 	return &MockServer{
 		mockTable: make(map[string]string),
 		authType:  MockTrust,
 	}
 }

 /*
 SetAuthType sets what kind of password authentication to require
 */
 func (m *MockServer) SetAuthType(auth MockAuth) {
 	m.authType = auth
 }

 /*
 SetTLSInfo sets the cert and key file and makes it possible for the server
 to support TLS.
 */
 func (m *MockServer) SetTLSInfo(certFile, keyFile string) error {
 	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
 	if err != nil {
 		return err
 	}
 	m.tlsConfig = &tls.Config{
 		Certificates: []tls.Certificate{cert},
 	}
 	return nil
 }

 /*
 SetForceTLS sets up the server to reject any non-TLS clients.
 */
 func (m *MockServer) SetForceTLS() {
 	m.forceTLS = true
 }

 /*
 Start listening for stuff.
 */
 func (m *MockServer) Start(port int) (err error) {
 	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
 		Port: port,
 	})

 	if err != nil {
 		return err
 	}

 	m.listener = listener
 	go m.acceptLoop()
 	return nil
 }

 /*
 Address returns the listen address in host:port format.
 */
 func (m *MockServer) Address() string {
 	return m.listener.Addr().String()
 }

 /*
 Stop stops the server listening for new connections.
 */
 func (m *MockServer) Stop() {
 	if m.listener != nil {
 		m.listener.Close()
 	}
 }

 /*
 acceptLoop sits and accepts new connections.
 */
 func (m *MockServer) acceptLoop() {
 	for {
 		conn, err := m.listener.Accept()
 		if err != nil {
 			return
 		}
 		go m.connectLoop(conn)
 	}
 }

 /*
 connectLoop responds to a new connection by handling the Postgres
 authentication and startup protocol.
 */
 func (m *MockServer) connectLoop(c net.Conn) {
 	defer c.Close()

 	startup, err := readMockMessage(c, true)
 	if err != nil {
 		log.Errorf("Error reading startup message: %s\n", err)
 		return
 	}

 	protoVersion, _ := startup.ReadInt32()

 	if protoVersion == sslMagicNumber {
 		// SSL startup attempt. Respond with "S" or "N"
 		if m.tlsConfig == nil {
 			c.Write([]byte{'N'})
 		} else {
 			// Respond with the appropriate byte and upgrade to TLS
 			c.Write([]byte{'S'})
 			c = tls.Server(c, m.tlsConfig)
 		}

 		// Look for a new startup packet now.
 		startup, err = readMockMessage(c, true)
 		if err != nil {
 			// This might happen if TLS handshake failed, which is a valid test case
 			return
 		}
 		protoVersion, _ = startup.ReadInt32()
 	} else if m.forceTLS {
 		// We will just close any connection that doesn't ask for TLS
 		return
 	}

 	if protoVersion != protocolVersion {
 		sendError(c, fmt.Sprintf("Invalid read protocol version %d\n", protoVersion))
 		return
 	}

 	var paramName, paramVal string
 	for {
 		paramName, _ = startup.ReadString()
 		if paramName == "" {
 			break
 		}
 		paramVal, _ = startup.ReadString()

 		if paramName == "user" {
 			if paramVal != mockUserName {
 				sendError(c, fmt.Sprintf("Invalid user name %s", paramVal))
 				return
 			}
 		}
 		if paramName == "database" {
 			if paramVal != mockDatabaseName {
 				sendError(c, fmt.Sprintf("Invalid database name %s\n", paramVal))
 				return
 			}
 		}
 	}

 	authOK := m.authLoop(c)

 	if authOK {
 		sendAuthResponse(c, 0)
 		sendReady(c)
 		m.readLoop(c)
 	}
 }

 /*
 authLoop runs the various Postgres authentication options.
 */
 func (m *MockServer) authLoop(c net.Conn) bool {
 	switch m.authType {
 	case MockTrust:
 		return true

 	case MockClear:
 		sendAuthResponse(c, 3)
 		return m.readPassword(c, mockPassword)

 	case MockMD5:
 		salt := make([]byte, 4)
 		rand.Read(salt)
 		out := NewServerOutputMessage(AuthenticationResponse)
 		out.WriteInt32(5)
 		out.WriteBytes(salt)
 		c.Write(out.Encode())
 		return m.readPassword(c, passwordMD5(mockUserName, mockPassword, salt))

 	default:
 		return false
 	}
 }

 func (m *MockServer) readPassword(c net.Conn, expected string) bool {
 	msg, _ := readMockMessage(c, false)
 	if msg.ServerType() != PasswordMessage {
 		sendError(c, "Expected PasswordMessage")
 		return false
 	}

 	pwd, _ := msg.ReadString()
 	if pwd == expected {
 		return true
 	}
 	sendError(c, "Invalid password")
 	return false
 }

 /*
 readLoop now reads and parses SQL commands until it's time to shut the
 connection down.
 */
 func (m *MockServer) readLoop(c net.Conn) {
 	state := mockIdle

 	for {
 		msg, err := readMockMessage(c, false)
 		if err != nil {
 			return
 		}

 		switch state {
 		case mockIdle:
 			m.readIdle(c, msg)
 		}
 	}
 }

 func (m *MockServer) readIdle(c net.Conn, msg *InputMessage) {
 	switch msg.ServerType() {
 	case Query:
 		sql, _ := msg.ReadString()
 		match := insertRE.FindStringSubmatch(sql)
 		if match != nil {
 			m.mockTable[match[1]] = match[2]
 			out := NewServerOutputMessage(CommandComplete)
 			out.WriteString("INSERT 1")
 			c.Write(out.Encode())
 			sendReady(c)

 		} else {
 			sendError(c, fmt.Sprintf("Invalid SQL \"%s\"", sql))
 			sendReady(c)
 		}

 	default:
 		sendError(c, fmt.Sprintf("Invalid message %s", msg.ServerType()))
 		sendReady(c)
 	}
 }

 func readMockMessage(c net.Conn, isStartup bool) (msg *InputMessage, err error) {
 	var hdr []byte
 	if isStartup {
 		hdr = make([]byte, 4)
 	} else {
 		hdr = make([]byte, 5)
 	}

 	_, err = io.ReadFull(c, hdr)
 	if err != nil {
 		return
 	}

 	hdrBuf := bytes.NewBuffer(hdr)
 	var msgType PgOutputType

 	if !isStartup {
 		var msgTypeVal byte
 		msgTypeVal, err = hdrBuf.ReadByte()
 		if err != nil {
 			return
 		}
 		msgType = PgOutputType(msgTypeVal)
 	}

 	var msgLen int32
 	err = binary.Read(hdrBuf, networkByteOrder, &msgLen)
 	if err != nil {
 		return
 	}

 	if msgLen < 4 {
 		err = fmt.Errorf("Invalid message length %d", msgLen)
 		return
 	}

 	bodBuf := make([]byte, msgLen-4)
 	_, err = io.ReadFull(c, bodBuf)
 	if err != nil {
 		return
 	}

 	msg = NewServerInputMessage(msgType, bodBuf)
 	return
 }

 func sendError(c net.Conn, msg string) {
 	out := NewServerOutputMessage(ErrorResponse)
 	out.WriteByte('S')
 	out.WriteString("FATAL")
 	out.WriteByte('M')
 	out.WriteString(msg)
 	out.WriteByte(0)
 	c.Write(out.Encode())
 }

 func sendAuthResponse(c net.Conn, code int32) {
 	out := NewServerOutputMessage(AuthenticationResponse)
 	out.WriteInt32(code)
 	c.Write(out.Encode())
 }

 func sendReady(c net.Conn) {
 	out := NewServerOutputMessage(ReadyForQuery)
 	out.WriteByte('I')
 	c.Write(out.Encode())
 }
	/*
	Copyright 2017 The Transicator Authors

	Licensed under the Apache License, Version 2.0 (the "License");
	you may not use this file except in compliance with the License.
	You may obtain a copy of the License at

	http://www.apache.org/licenses/LICENSE-2.0

	Unless required by applicable law or agreed to in writing, software
	distributed under the License is distributed on an "AS IS" BASIS,
	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	See the License for the specific language governing permissions and
	limitations under the License.
	*/
	package pgclient

	import (
	"bytes"
	"crypto/rand"
	"crypto/tls"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"regexp"

	log "github.com/Sirupsen/logrus"
	)

	const (
	mockUserName = "mock"
	mockPassword = "mocketty"
	mockDatabaseName = "turtle"
	)

	type mockState int

	const (
	mockIdle mockState = 1
	)

	// MockAuth specifies what type of authentication the server supports
	type MockAuth int

	// Different auth types
	const (
	MockTrust MockAuth = 0
	MockClear MockAuth = 1
	MockMD5 MockAuth = 2
	)

	var insertRE = regexp.MustCompile("insert into mock values \\('([\\w]+)', '([\\w]+)'\\)")

	/*
	A MockServer is a server that implements a little bit of the Postgres wire
	protocol. We can use it for testing of the wire protocol client. In particular,
	we can use it to test the myriad of password authentication and TLS options
	without having to start and stop a real Postgres server in the test suite.
	*/
	type MockServer struct {
	listener net.Listener
	mockTable map[string]string
	authType MockAuth
	tlsConfig *tls.Config
	forceTLS bool
	}

	/*
	NewMockServer starts a new server in the current process, listening on the
	specified port.
	*/
	func NewMockServer() *MockServer {
	return &MockServer{
	mockTable: make(map[string]string),
	authType: MockTrust,
	}
	}

	/*
	SetAuthType sets what kind of password authentication to require
	*/
	func (m *MockServer) SetAuthType(auth MockAuth) {
	m.authType = auth
	}

	/*
	SetTLSInfo sets the cert and key file and makes it possible for the server
	to support TLS.
	*/
	func (m *MockServer) SetTLSInfo(certFile, keyFile string) error {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
	return err
	}
	m.tlsConfig = &tls.Config{
	Certificates: []tls.Certificate{cert},
	}
	return nil
	}

	/*
	SetForceTLS sets up the server to reject any non-TLS clients.
	*/
	func (m *MockServer) SetForceTLS() {
	m.forceTLS = true
	}

	/*
	Start listening for stuff.
	*/
	func (m *MockServer) Start(port int) (err error) {
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
	Port: port,
	})

	if err != nil {
	return err
	}

	m.listener = listener
	go m.acceptLoop()
	return nil
	}

	/*
	Address returns the listen address in host:port format.
	*/
	func (m *MockServer) Address() string {
	return m.listener.Addr().String()
	}

	/*
	Stop stops the server listening for new connections.
	*/
	func (m *MockServer) Stop() {
	if m.listener != nil {
	m.listener.Close()
	}
	}

	/*
	acceptLoop sits and accepts new connections.
	*/
	func (m *MockServer) acceptLoop() {
	for {
	conn, err := m.listener.Accept()
	if err != nil {
	return
	}
	go m.connectLoop(conn)
	}
	}

	/*
	connectLoop responds to a new connection by handling the Postgres
	authentication and startup protocol.
	*/
	func (m *MockServer) connectLoop(c net.Conn) {
	defer c.Close()

	startup, err := readMockMessage(c, true)
	if err != nil {
	log.Errorf("Error reading startup message: %s\n", err)
	return
	}

	protoVersion, _ := startup.ReadInt32()

	if protoVersion == sslMagicNumber {
	// SSL startup attempt. Respond with "S" or "N"
	if m.tlsConfig == nil {
	c.Write([]byte{'N'})
	} else {
	// Respond with the appropriate byte and upgrade to TLS
	c.Write([]byte{'S'})
	c = tls.Server(c, m.tlsConfig)
	}

	// Look for a new startup packet now.
	startup, err = readMockMessage(c, true)
	if err != nil {
	// This might happen if TLS handshake failed, which is a valid test case
	return
	}
	protoVersion, _ = startup.ReadInt32()
	} else if m.forceTLS {
	// We will just close any connection that doesn't ask for TLS
	return
	}

	if protoVersion != protocolVersion {
	sendError(c, fmt.Sprintf("Invalid read protocol version %d\n", protoVersion))
	return
	}

	var paramName, paramVal string
	for {
	paramName, _ = startup.ReadString()
	if paramName == "" {
	break
	}
	paramVal, _ = startup.ReadString()

	if paramName == "user" {
	if paramVal != mockUserName {
	sendError(c, fmt.Sprintf("Invalid user name %s", paramVal))
	return
	}
	}
	if paramName == "database" {
	if paramVal != mockDatabaseName {
	sendError(c, fmt.Sprintf("Invalid database name %s\n", paramVal))
	return
	}
	}
	}

	authOK := m.authLoop(c)

	if authOK {
	sendAuthResponse(c, 0)
	sendReady(c)
	m.readLoop(c)
	}
	}

	/*
	authLoop runs the various Postgres authentication options.
	*/
	func (m *MockServer) authLoop(c net.Conn) bool {
	switch m.authType {
	case MockTrust:
	return true

	case MockClear:
	sendAuthResponse(c, 3)
	return m.readPassword(c, mockPassword)

	case MockMD5:
	salt := make([]byte, 4)
	rand.Read(salt)
	out := NewServerOutputMessage(AuthenticationResponse)
	out.WriteInt32(5)
	out.WriteBytes(salt)
	c.Write(out.Encode())
	return m.readPassword(c, passwordMD5(mockUserName, mockPassword, salt))

	default:
	return false
	}
	}

	func (m *MockServer) readPassword(c net.Conn, expected string) bool {
	msg, _ := readMockMessage(c, false)
	if msg.ServerType() != PasswordMessage {
	sendError(c, "Expected PasswordMessage")
	return false
	}

	pwd, _ := msg.ReadString()
	if pwd == expected {
	return true
	}
	sendError(c, "Invalid password")
	return false
	}

	/*
	readLoop now reads and parses SQL commands until it's time to shut the
	connection down.
	*/
	func (m *MockServer) readLoop(c net.Conn) {
	state := mockIdle

	for {
	msg, err := readMockMessage(c, false)
	if err != nil {
	return
	}

	switch state {
	case mockIdle:
	m.readIdle(c, msg)
	}
	}
	}

	func (m MockServer) readIdle(c net.Conn, msg InputMessage) {
	switch msg.ServerType() {
	case Query:
	sql, _ := msg.ReadString()
	match := insertRE.FindStringSubmatch(sql)
	if match != nil {
	m.mockTable[match[1]] = match[2]
	out := NewServerOutputMessage(CommandComplete)
	out.WriteString("INSERT 1")
	c.Write(out.Encode())
	sendReady(c)

	} else {
	sendError(c, fmt.Sprintf("Invalid SQL \"%s\"", sql))
	sendReady(c)
	}

	default:
	sendError(c, fmt.Sprintf("Invalid message %s", msg.ServerType()))
	sendReady(c)
	}
	}

	func readMockMessage(c net.Conn, isStartup bool) (msg *InputMessage, err error) {
	var hdr []byte
	if isStartup {
	hdr = make([]byte, 4)
	} else {
	hdr = make([]byte, 5)
	}

	_, err = io.ReadFull(c, hdr)
	if err != nil {
	return
	}

	hdrBuf := bytes.NewBuffer(hdr)
	var msgType PgOutputType

	if !isStartup {
	var msgTypeVal byte
	msgTypeVal, err = hdrBuf.ReadByte()
	if err != nil {
	return
	}
	msgType = PgOutputType(msgTypeVal)
	}

	var msgLen int32
	err = binary.Read(hdrBuf, networkByteOrder, &msgLen)
	if err != nil {
	return
	}

	if msgLen < 4 {
	err = fmt.Errorf("Invalid message length %d", msgLen)
	return
	}

	bodBuf := make([]byte, msgLen-4)
	_, err = io.ReadFull(c, bodBuf)
	if err != nil {
	return
	}

	msg = NewServerInputMessage(msgType, bodBuf)
	return
	}

	func sendError(c net.Conn, msg string) {
	out := NewServerOutputMessage(ErrorResponse)
	out.WriteByte('S')
	out.WriteString("FATAL")
	out.WriteByte('M')
	out.WriteString(msg)
	out.WriteByte(0)
	c.Write(out.Encode())
	}

	func sendAuthResponse(c net.Conn, code int32) {
	out := NewServerOutputMessage(AuthenticationResponse)
	out.WriteInt32(code)
	c.Write(out.Encode())
	}

	func sendReady(c net.Conn) {
	out := NewServerOutputMessage(ReadyForQuery)
	out.WriteByte('I')
	c.Write(out.Encode())
	}