dolthub--doltgresql/testing/dumps/intercept.go

// Copyright 2025 Dolthub, Inc.
//
// 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 dumps

import (
	"fmt"
	"net"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/dolthub/go-mysql-server/server"
	"github.com/dolthub/go-mysql-server/sql"
	"github.com/jackc/pgx/v5/pgproto3"
	"github.com/stretchr/testify/require"
)

// ImportQueryError contains both a query and its associated error.
type ImportQueryError struct {
	Query string
	Error string
}

// InterceptArgs are the arguments that are passed to InterceptImportMessages.
type InterceptArgs struct {
	DoltgresPort      int
	SkippedQueries    []string
	BreakpointQueries []string
	TriggerBreakpoint func(string)
}

// passthroughArguments are the arguments that are passed to createPassthrough.
type passthroughArguments struct {
	qeChan               chan ImportQueryError
	terminate            *sync.WaitGroup
	psqlConnBackend      *pgproto3.Backend
	doltgresConnFrontend *pgproto3.Frontend
	triggerBreakpoint    func(string)
	skippedQueries       []string
	breakpointQueries    []string
}

// InterceptImportMessages sits between PSQL and Doltgres, returning all error messages that are encountered. As we rely
// on PSQL to handle the import process, we normally wouldn't be able to associate error messages with queries, as this
// information is not returned by PSQL itself. Therefore, we create our own connection to Doltgres, and a server that
// PSQL listens to. We then forward everything from PSQL to Doltgres, while inspecting the messages as they come and go.
func InterceptImportMessages(t *testing.T, args InterceptArgs) (int, chan ImportQueryError) {
	psqlPort, err := sql.GetEmptyPort()
	require.NoError(t, err)
	qeChan := make(chan ImportQueryError)
	listener, err := server.NewListener("tcp", fmt.Sprintf("127.0.0.1:%d", psqlPort), "")
	if err != nil {
		t.Fatal(err)
		return psqlPort, qeChan
	}
	timer := time.NewTimer(5 * time.Second)
	timer.Stop()
	go func() {
		<-timer.C
		_ = listener.Close()
	}()

	go func() {
		for {
			psqlConn, err := listener.Accept()
			if err != nil {
				return
			}
			timer.Stop()
			terminate := &sync.WaitGroup{}
			terminate.Add(1)
			psqlConnBackend := pgproto3.NewBackend(psqlConn, psqlConn)
			doltgresConn, err := (&net.Dialer{}).Dial("tcp", fmt.Sprintf("127.0.0.1:%d", args.DoltgresPort))
			if err != nil {
				fmt.Println(err)
				return
			}
			doltgresConnFrontend := pgproto3.NewFrontend(doltgresConn, doltgresConn)

			if err = handleStartup(t, psqlConnBackend, doltgresConnFrontend, psqlConn); err != nil {
				fmt.Println(err)
				return
			}
			createPassthrough(passthroughArguments{
				qeChan:               qeChan,
				terminate:            terminate,
				psqlConnBackend:      psqlConnBackend,
				doltgresConnFrontend: doltgresConnFrontend,
				triggerBreakpoint:    args.TriggerBreakpoint,
				skippedQueries:       args.SkippedQueries,
				breakpointQueries:    args.BreakpointQueries,
			})
			terminate.Wait()
			_ = psqlConn.Close()
			_ = doltgresConn.Close()
			timer.Reset(5 * time.Second)
		}
	}()
	return psqlPort, qeChan
}

// handleStartup handles the startup messages.
func handleStartup(t *testing.T, psqlConnBackend *pgproto3.Backend, doltgresConnFrontend *pgproto3.Frontend, clientConn net.Conn) error {
StartupLoop:
	for {
		startupMessage, err := psqlConnBackend.ReceiveStartupMessage()
		if err != nil {
			return err
		}
		switch startupMessage := startupMessage.(type) {
		case *pgproto3.SSLRequest:
			if _, err = clientConn.Write([]byte{'N'}); err != nil {
				return err
			}
		case *pgproto3.StartupMessage:
			doltgresConnFrontend.Send(startupMessage)
			if err = doltgresConnFrontend.Flush(); err != nil {
				return err
			}
			response, err := doltgresConnFrontend.Receive()
			if err != nil {
				return err
			}
			if err = setAuthType(psqlConnBackend, response); err != nil {
				return err
			}
			psqlConnBackend.Send(response)
			if err = psqlConnBackend.Flush(); err != nil {
				return err
			}
			break StartupLoop
		case *pgproto3.GSSEncRequest:
			// we don't support GSSAPI
			_, err = clientConn.Write([]byte("N"))
			if err != nil {
				return err
			}
		default:
			t.Fatalf("unexpected startup message: %v", startupMessage)
		}
	}
	return nil
}

// setAuthType sets the client's authentication type depending on the message received from the server. This is
// necessary, as the client needs the proper context to know how to parse the returned messages.
func setAuthType(clientConnBackend *pgproto3.Backend, message pgproto3.BackendMessage) error {
	switch message.(type) {
	case *pgproto3.AuthenticationOk:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeOk)
	case *pgproto3.AuthenticationCleartextPassword:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeCleartextPassword)
	case *pgproto3.AuthenticationMD5Password:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeMD5Password)
	case *pgproto3.AuthenticationGSS:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeGSS)
	case *pgproto3.AuthenticationGSSContinue:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeGSSCont)
	case *pgproto3.AuthenticationSASL:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeSASL)
	case *pgproto3.AuthenticationSASLContinue:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeSASLContinue)
	case *pgproto3.AuthenticationSASLFinal:
		return clientConnBackend.SetAuthType(pgproto3.AuthTypeSASLFinal)
	default:
		return nil
	}
}

// createPassthrough creates the go routines that will read from and write to the connections.
func createPassthrough(args passthroughArguments) {
	lastQuery := ""
	writeMutex := &sync.Mutex{}
	go func() {
		defer args.terminate.Done()
		for {
			psqlMessage, err := args.psqlConnBackend.Receive()
			if err != nil {
				errStr := err.Error()
				if errStr != "unexpected EOF" && !strings.HasSuffix(errStr, "use of closed network connection") {
					fmt.Println(err)
				}
				return
			}
			switch msg := psqlMessage.(type) {
			case *pgproto3.Query:
				writeMutex.Lock()
				if len(lastQuery) == 0 {
					lastQuery = msg.String
				}
				writeMutex.Unlock()
				for _, query := range args.skippedQueries {
					if strings.HasPrefix(msg.String, query) {
						// An empty query allows for the proper response messages to be sent.
						msg.String = ";"
						break
					}
				}
				for _, query := range args.breakpointQueries {
					if strings.HasPrefix(msg.String, query) {
						args.triggerBreakpoint(msg.String)
						break
					}
				}
			case *pgproto3.Terminate:
				return
			}
			args.doltgresConnFrontend.Send(psqlMessage)
			if err = args.doltgresConnFrontend.Flush(); err != nil {
				errStr := err.Error()
				if errStr != "unexpected EOF" && !strings.HasSuffix(errStr, "use of closed network connection") {
					fmt.Println(err)
				}
				return
			}
		}
	}()
	go func() {
		for {
			doltgresMessage, err := args.doltgresConnFrontend.Receive()
			if err != nil {
				errStr := err.Error()
				if errStr != "unexpected EOF" &&
					!strings.HasSuffix(errStr, "use of closed network connection") &&
					!strings.HasSuffix(errStr, "An existing connection was forcibly closed by the remote host.") {
					fmt.Println(err)
				}
				return
			}
			switch msg := doltgresMessage.(type) {
			case *pgproto3.ErrorResponse:
				writeMutex.Lock()
				if len(lastQuery) == 0 {
					args.qeChan <- ImportQueryError{
						Query: "UNKNOWN QUERY HAS ERRORED",
						Error: msg.Message,
					}
				} else {
					args.qeChan <- ImportQueryError{
						Query: lastQuery,
						Error: msg.Message,
					}
				}
				writeMutex.Unlock()
			case *pgproto3.ReadyForQuery:
				writeMutex.Lock()
				lastQuery = ""
				writeMutex.Unlock()
			default:
				if err = setAuthType(args.psqlConnBackend, doltgresMessage); err != nil {
					fmt.Println(err)
					return
				}
			}
			args.psqlConnBackend.Send(doltgresMessage)
			if err = args.psqlConnBackend.Flush(); err != nil {
				errStr := err.Error()
				if errStr != "unexpected EOF" &&
					!strings.HasSuffix(errStr, "use of closed network connection") &&
					!strings.HasSuffix(errStr, "An existing connection was forcibly closed by the remote host.") {
					fmt.Println(err)
				}
				return
			}
		}
	}()
}