16306: Merge branch 'master'
[arvados.git] / lib / controller / localdb / docker_test.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package localdb
6
7 import (
8         "io"
9         "net"
10         "strings"
11
12         "git.arvados.org/arvados.git/sdk/go/arvados"
13         check "gopkg.in/check.v1"
14 )
15
16 type pgproxy struct {
17         net.Listener
18 }
19
20 // newPgProxy sets up a TCP proxy, listening on all interfaces, that
21 // forwards all connections to the cluster's PostgreSQL server. This
22 // allows the caller to run a docker container that can connect to a
23 // postgresql instance that listens on the test host's loopback
24 // interface.
25 //
26 // Caller is responsible for calling Close() on the returned pgproxy.
27 func newPgProxy(c *check.C, cluster *arvados.Cluster) *pgproxy {
28         host := cluster.PostgreSQL.Connection["host"]
29         if host == "" {
30                 host = "localhost"
31         }
32         port := cluster.PostgreSQL.Connection["port"]
33         if port == "" {
34                 port = "5432"
35         }
36         target := net.JoinHostPort(host, port)
37
38         ln, err := net.Listen("tcp", ":")
39         c.Assert(err, check.IsNil)
40         go func() {
41                 for {
42                         downstream, err := ln.Accept()
43                         if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
44                                 return
45                         }
46                         c.Assert(err, check.IsNil)
47                         go func() {
48                                 c.Logf("pgproxy accepted connection from %s", downstream.RemoteAddr().String())
49                                 defer downstream.Close()
50                                 upstream, err := net.Dial("tcp", target)
51                                 if err != nil {
52                                         c.Logf("net.Dial(%q): %s", target, err)
53                                         return
54                                 }
55                                 defer upstream.Close()
56                                 go io.Copy(downstream, upstream)
57                                 io.Copy(upstream, downstream)
58                         }()
59                 }
60         }()
61         c.Logf("pgproxy listening at %s", ln.Addr().String())
62         return &pgproxy{Listener: ln}
63 }
64
65 func (proxy *pgproxy) Port() string {
66         _, port, _ := net.SplitHostPort(proxy.Addr().String())
67         return port
68 }