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 }