postgres协议学习

type PgProxyConnection struct {
	clientConn    net.Conn           // tcp connection between client and proxy
	pgConn        net.Conn           // tcp connection between proxy and postgres
	clientControl *pgproto3.Backend  // clientConn controller
	pgControl     *pgproto3.Frontend // pgConn controller
	TlsOn         bool               // is tls on
	ctx           context.Context
}

func (pc *PgProxyConnection) startUp() error {
	if pc.TlsOn {
		crt, err := tls.LoadX509KeyPair("server.crt", "server.key")
		if err != nil {
			return err
		}
		cfg := &tls.Config{
			Certificates:       []tls.Certificate{crt},
			InsecureSkipVerify: true,
			MinVersion:         tls.VersionTLS12,
			MaxVersion:         tls.VersionTLS12,
		}
		// upgrade to tls connection
		pc.clientConn = tls.Server(pc.clientConn, cfg)
		pc.clientControl = pgproto3.NewBackend(pgproto3.NewChunkReader(pc.clientConn), pc.clientConn)
	}
	startUpMsg, err := pc.clientControl.ReceiveStartupMessage()
	if err != nil {
		return err
	}

	switch msg := startUpMsg.(type) {
	case *pgproto3.StartupMessage:
		return pc.handShake(msg)
	case *pgproto3.SSLRequest:
		// set TLS option on
		pc.TlsOn = true
		// ACK for ssl request
		_, err = pc.clientConn.Write([]byte{'S'})
		if err != nil {
			return err
		}
		return pc.startUp()
	default:
		return fmt.Errorf("unknown startup message: %#v", startUpMsg)
	}
}

func (pc *PgProxyConnection) handShake(msg *pgproto3.StartupMessage) error {
	host := msg.Parameters["host"]
	port := msg.Parameters["port"]
	var err error
	dialer := net.Dialer{KeepAlive: 5 * time.Minute}
	if pc.pgConn, err = dialer.Dial("tcp", fmt.Sprintf("%s:%s", host, port)); err != nil {
		return err
	}
	if pc.TlsOn {
		sslMode, _ := config.Get("TCP.SslMode").(string)
		rootCrt, _ := config.Get("TCP.RootCrt").(string)
		serverCrt, _ := config.Get("TCP.ServerCrt").(string)
		serverKey, _ := config.Get("TCP.ServerKey").(string)
		var cfg *tls.Config
		if cfg, err = configTLS(sslMode, rootCrt, serverCrt, serverKey, host); err != nil {
			return err
		}
		if pc.pgConn, err = startTLS(pc.pgConn, cfg); err != nil {
			return err
		}
	}
	cr, err := chunkreader.NewConfig(pc.pgConn, chunkreader.Config{MinBufLen: 8192})
	if err != nil {
		return err
	}
	pc.pgControl = pgproto3.NewFrontend(cr, pc.pgConn)

	if _, err := pc.pgConn.Write(msg.Encode(nil)); err != nil {
		pc.pgConn.Close()
		return err
	}

	for {
		pdbStartUpMsg, err := pc.pgControl.Receive()
		if err != nil {
			return err
		}
		switch parsedMsg := pdbStartUpMsg.(type) {
		case *pgproto3.ReadyForQuery:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
			return nil
		case *pgproto3.ParameterStatus:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		case *pgproto3.ErrorResponse:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
			return errors.New(parsedMsg.Message)
		case *pgproto3.NoticeResponse:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		case *pgproto3.NotificationResponse:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		case *pgproto3.BackendKeyData:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		case *pgproto3.AuthenticationOk:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		case *pgproto3.AuthenticationCleartextPassword:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		case *pgproto3.AuthenticationMD5Password:
            // 对密码的md5生成规则为:
            // concat('md5', md5(concat(md5(concat(password, username)), random-salt)))
			_, err = pc.clientConn.Write(parsedMsg.Encode(nil))
			if err != nil {
				return err
			}
			md5pw, err := pc.clientControl.Receive()
			if err != nil {
				return err
			}
			_, err = pc.pgConn.Write(md5pw.Encode(nil))
			if err != nil {
				return err
			}
		case *pgproto3.AuthenticationSASL:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		case *pgproto3.AuthenticationGSS:
			_, _ = pc.clientConn.Write(parsedMsg.Encode(nil))
		default:
			return errors.New("unknown connect message")
		}
	}
}