<commit_before>package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/signal"
	"os/user"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/packethost/packngo"
	log "github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/terminal"
)

const (
	packetDefaultZone    = "ams1"
	packetDefaultMachine = "baremetal_0"
	packetBaseURL        = "PACKET_BASE_URL"
	packetZoneVar        = "PACKET_ZONE"
	packetMachineVar     = "PACKET_MACHINE"
	packetAPIKeyVar      = "PACKET_API_KEY"
	packetProjectIDVar   = "PACKET_PROJECT_ID"
	packetHostnameVar    = "PACKET_HOSTNAME"
	packetNameVar        = "PACKET_NAME"
)

var (
	packetDefaultHostname = "linuxkit"
)

func init() {
	// Prefix host name with username
	if u, err := user.Current(); err == nil {
		packetDefaultHostname = u.Username + "-" + packetDefaultHostname
	}
}

// Process the run arguments and execute run
func runPacket(args []string) {
	flags := flag.NewFlagSet("packet", flag.ExitOnError)
	invoked := filepath.Base(os.Args[0])
	flags.Usage = func() {
		fmt.Printf("USAGE: %s run packet [options] [name]\n\n", invoked)
		fmt.Printf("Options:\n\n")
		flags.PrintDefaults()
	}
	baseURLFlag := flags.String("base-url", "", "Base URL that the kernel and initrd are served from (or "+packetBaseURL+")")
	zoneFlag := flags.String("zone", packetDefaultZone, "Packet Zone (or "+packetZoneVar+")")
	machineFlag := flags.String("machine", packetDefaultMachine, "Packet Machine Type (or "+packetMachineVar+")")
	apiKeyFlag := flags.String("api-key", "", "Packet API key (or "+packetAPIKeyVar+")")
	projectFlag := flags.String("project-id", "", "Packet Project ID (or "+packetProjectIDVar+")")
	hostNameFlag := flags.String("hostname", packetDefaultHostname, "Hostname of new instance (or "+packetHostnameVar+")")
	nameFlag := flags.String("img-name", "", "Overrides the prefix used to identify the files. Defaults to [name] (or "+packetNameVar+")")
	alwaysPXE := flags.Bool("always-pxe", true, "Reboot from PXE every time.")
	serveFlag := flags.String("serve", "", "Serve local files via the http port specified, e.g. ':8080'.")
	consoleFlag := flags.Bool("console", true, "Provide interactive access on the console.")
	keepFlag := flags.Bool("keep", false, "Keep the machine after exiting/poweroff.")
	if err := flags.Parse(args); err != nil {
		log.Fatal("Unable to parse args")
	}

	remArgs := flags.Args()
	prefix := "packet"
	if len(remArgs) > 0 {
		prefix = remArgs[0]
	}

	url := getStringValue(packetBaseURL, *baseURLFlag, "")
	if url == "" {
		log.Fatal("Need to specify a value for --base-url where the images are hosted. This URL should contain <url>/%s-kernel and <url>/%s-initrd.img")
	}
	facility := getStringValue(packetZoneVar, *zoneFlag, "")
	plan := getStringValue(packetMachineVar, *machineFlag, defaultMachine)
	apiKey := getStringValue(packetAPIKeyVar, *apiKeyFlag, "")
	if apiKey == "" {
		log.Fatal("Must specify a Packet.net API key with --api-key")
	}
	projectID := getStringValue(packetProjectIDVar, *projectFlag, "")
	if projectID == "" {
		log.Fatal("Must specify a Packet.net Project ID with --project-id")
	}
	hostname := getStringValue(packetHostnameVar, *hostNameFlag, "")
	name := getStringValue(packetNameVar, *nameFlag, prefix)
	osType := "custom_ipxe"
	billing := "hourly"

	if !*keepFlag && !*consoleFlag {
		log.Fatalf("Combination of keep=%t and console=%t makes little sense", *keepFlag, *consoleFlag)
	}

	// Read kernel command line
	var cmdline string
	if c, err := ioutil.ReadFile(prefix + "-cmdline"); err != nil {
		log.Fatalf("Cannot open cmdline file: %v", err)
	} else {
		cmdline = string(c)
	}

	// Serve files with a local http server
	var httpServer *http.Server
	if *serveFlag != "" {
		fs := serveFiles{[]string{fmt.Sprintf("%s-kernel", name), fmt.Sprintf("%s-initrd.img", name)}}
		httpServer = &http.Server{Addr: ":8080", Handler: http.FileServer(fs)}
		go func() {
			log.Infof("Listening on http://%s\n", *serveFlag)
			if err := httpServer.ListenAndServe(); err != nil {
				log.Infof("http server exited with: %v", err)
			}
		}()
	}

	// Build the iPXE script
	// Note, we *append* the <prefix>-cmdline. iXPE booting will
	// need the first set of "kernel-params" and we don't want to
	// require these to be added to every YAML file.
	userData := "#!ipxe\n\n"
	userData += "dhcp\n"
	userData += fmt.Sprintf("set base-url %s\n", url)
	userData += fmt.Sprintf("set kernel-params ip=dhcp nomodeset ro serial console=ttyS1,115200 %s\n", cmdline)
	userData += fmt.Sprintf("kernel ${base-url}/%s-kernel ${kernel-params}\n", name)
	userData += fmt.Sprintf("initrd ${base-url}/%s-initrd.img\n", name)
	userData += "boot"
	log.Debugf("Using userData of:\n%s\n", userData)

	// Make sure the URL works
	initrdURL := fmt.Sprintf("%s/%s-initrd.img", url, name)
	kernelURL := fmt.Sprintf("%s/%s-kernel", url, name)
	validateHTTPURL(kernelURL)
	validateHTTPURL(initrdURL)

	client := packngo.NewClient("", apiKey, nil)
	tags := []string{}
	req := packngo.DeviceCreateRequest{
		HostName:     hostname,
		Plan:         plan,
		Facility:     facility,
		OS:           osType,
		BillingCycle: billing,
		ProjectID:    projectID,
		UserData:     userData,
		Tags:         tags,
		AlwaysPXE:    *alwaysPXE,
	}
	dev, _, err := client.Devices.Create(&req)
	if err != nil {
		log.Fatal(err)
	}
	b, err := json.MarshalIndent(dev, "", "    ")
	if err != nil {
		log.Fatal(err)
	}
	// log response json if in verbose mode
	log.Debugf("%s\n", string(b))

	sshHost := "sos." + dev.Facility.Code + ".packet.net"
	if *consoleFlag {
		// Connect to the serial console
		if err := sshSOS(dev.ID, sshHost); err != nil {
			log.Fatal(err)
		}
	} else {
		log.Printf("Machine booting")
		log.Printf("Access the console with: ssh %s@%s", dev.ID, sshHost)

		// if the serve option is present, wait till 'ctrl-c' is hit.
		// Otherwise we wouldn't serve the files
		if *serveFlag != "" {
			stop := make(chan os.Signal, 1)
			signal.Notify(stop, os.Interrupt)
			log.Printf("Hit ctrl-c to stop http server")
			<-stop
		}
	}

	// Stop the http server before exiting
	if *serveFlag != "" {
		log.Printf("Shutting down http server...")
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		httpServer.Shutdown(ctx)
	}

	if !*keepFlag {
		if _, err := client.Devices.Delete(dev.ID); err != nil {
			log.Fatalf("Unable to delete device: %v", err)
		}
	}

}

// validateHTTPURL does a sanity check that a URL returns a 200 or 300 response
func validateHTTPURL(url string) {
	log.Printf("Validating URL: %s", url)
	resp, err := http.Head(url)
	if err != nil {
		log.Fatal(err)
	}
	if resp.StatusCode >= 400 {
		log.Fatal("Got a non 200- or 300- HTTP response code: %s", resp)
	}
	log.Printf("OK: %d response code", resp.StatusCode)
}

func sshSOS(user, host string) error {
	log.Printf("console: ssh %s@%s", user, host)

	hostKey, err := sshHostKey(host)
	if err != nil {
		return fmt.Errorf("Host key not found. Maybe need to add it? %v", err)
	}

	sshConfig := &ssh.ClientConfig{
		User:            user,
		HostKeyCallback: ssh.FixedHostKey(hostKey),
		Auth: []ssh.AuthMethod{
			sshAgent(),
		},
	}

	c, err := ssh.Dial("tcp", host+":22", sshConfig)
	if err != nil {
		return fmt.Errorf("Failed to dial: %s", err)
	}

	s, err := c.NewSession()
	if err != nil {
		return fmt.Errorf("Failed to create session: %v", err)
	}
	defer s.Close()

	s.Stdout = os.Stdout
	s.Stderr = os.Stderr
	s.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO:  0,
		ssh.IGNCR: 1,
	}

	width, height, err := terminal.GetSize(0)
	if err != nil {
		log.Warningf("Error getting terminal size. Ignored. %v", err)
		width = 80
		height = 40
	}
	if err := s.RequestPty("vt100", width, height, modes); err != nil {
		return fmt.Errorf("Request for PTY failed: %v", err)
	}
	oldState, err := terminal.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return err
	}
	defer terminal.Restore(0, oldState)

	// Start remote shell
	if err := s.Shell(); err != nil {
		return fmt.Errorf("Failed to start shell: %v", err)
	}

	s.Wait()
	return nil
}

// Get a ssh-agent AuthMethod
func sshAgent() ssh.AuthMethod {
	sshAgent, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		log.Fatalf("Failed to dial ssh-agent: %v", err)
	}
	return ssh.PublicKeysCallback(agent.NewClient(sshAgent).Signers)
}

// This function returns the host key for a given host (the SOS server).
// If it can't be found, it errors
func sshHostKey(host string) (ssh.PublicKey, error) {
	f, err := os.Open(filepath.Join(os.Getenv("HOME"), ".ssh", "known_hosts"))
	if err != nil {
		return nil, fmt.Errorf("Can't open know_hosts file: %v", err)
	}
	defer f.Close()

	s := bufio.NewScanner(f)

	var hostKey ssh.PublicKey
	for s.Scan() {
		fields := strings.Split(s.Text(), " ")
		if len(fields) != 3 {
			continue
		}
		if strings.Contains(fields[0], host) {
			var err error
			hostKey, _, _, _, err = ssh.ParseAuthorizedKey(s.Bytes())
			if err != nil {
				return nil, fmt.Errorf("Error parsing %q: %v", fields[2], err)
			}
			break
		}
	}

	if hostKey == nil {
		return nil, fmt.Errorf("No hostkey for %s", host)
	}
	return hostKey, nil
}

// This implements a http.FileSystem which only responds to specific files.
type serveFiles struct {
	files []string
}

// Open implements the Open method for the serveFiles FileSystem
// implementation.
// It converts both the name from the URL and the files provided in
// the serveFiles structure into cleaned, absolute filesystem path and
// only returns the file if the requested name matches one of the
// files in the list.
func (fs serveFiles) Open(name string) (http.File, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	name = filepath.Join(cwd, filepath.FromSlash(path.Clean("/"+name)))
	for _, fn := range fs.files {
		fn = filepath.Join(cwd, filepath.FromSlash(path.Clean("/"+fn)))
		if name == fn {
			f, err := os.Open(fn)
			if err != nil {
				return nil, err
			}
			log.Infof("Serving: %s", fn)
			return f, nil
		}
	}
	return nil, fmt.Errorf("File %s not found", name)
}
<commit_msg>cmd: Unify Logging for packet backend<commit_after>package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/signal"
	"os/user"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/packethost/packngo"
	log "github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/crypto/ssh/terminal"
)

const (
	packetDefaultZone    = "ams1"
	packetDefaultMachine = "baremetal_0"
	packetBaseURL        = "PACKET_BASE_URL"
	packetZoneVar        = "PACKET_ZONE"
	packetMachineVar     = "PACKET_MACHINE"
	packetAPIKeyVar      = "PACKET_API_KEY"
	packetProjectIDVar   = "PACKET_PROJECT_ID"
	packetHostnameVar    = "PACKET_HOSTNAME"
	packetNameVar        = "PACKET_NAME"
)

var (
	packetDefaultHostname = "linuxkit"
)

func init() {
	// Prefix host name with username
	if u, err := user.Current(); err == nil {
		packetDefaultHostname = u.Username + "-" + packetDefaultHostname
	}
}

// Process the run arguments and execute run
func runPacket(args []string) {
	flags := flag.NewFlagSet("packet", flag.ExitOnError)
	invoked := filepath.Base(os.Args[0])
	flags.Usage = func() {
		fmt.Printf("USAGE: %s run packet [options] [name]\n\n", invoked)
		fmt.Printf("Options:\n\n")
		flags.PrintDefaults()
	}
	baseURLFlag := flags.String("base-url", "", "Base URL that the kernel and initrd are served from (or "+packetBaseURL+")")
	zoneFlag := flags.String("zone", packetDefaultZone, "Packet Zone (or "+packetZoneVar+")")
	machineFlag := flags.String("machine", packetDefaultMachine, "Packet Machine Type (or "+packetMachineVar+")")
	apiKeyFlag := flags.String("api-key", "", "Packet API key (or "+packetAPIKeyVar+")")
	projectFlag := flags.String("project-id", "", "Packet Project ID (or "+packetProjectIDVar+")")
	hostNameFlag := flags.String("hostname", packetDefaultHostname, "Hostname of new instance (or "+packetHostnameVar+")")
	nameFlag := flags.String("img-name", "", "Overrides the prefix used to identify the files. Defaults to [name] (or "+packetNameVar+")")
	alwaysPXE := flags.Bool("always-pxe", true, "Reboot from PXE every time.")
	serveFlag := flags.String("serve", "", "Serve local files via the http port specified, e.g. ':8080'.")
	consoleFlag := flags.Bool("console", true, "Provide interactive access on the console.")
	keepFlag := flags.Bool("keep", false, "Keep the machine after exiting/poweroff.")
	if err := flags.Parse(args); err != nil {
		log.Fatal("Unable to parse args")
	}

	remArgs := flags.Args()
	prefix := "packet"
	if len(remArgs) > 0 {
		prefix = remArgs[0]
	}

	url := getStringValue(packetBaseURL, *baseURLFlag, "")
	if url == "" {
		log.Fatal("Need to specify a value for --base-url where the images are hosted. This URL should contain <url>/%s-kernel and <url>/%s-initrd.img")
	}
	facility := getStringValue(packetZoneVar, *zoneFlag, "")
	plan := getStringValue(packetMachineVar, *machineFlag, defaultMachine)
	apiKey := getStringValue(packetAPIKeyVar, *apiKeyFlag, "")
	if apiKey == "" {
		log.Fatal("Must specify a Packet.net API key with --api-key")
	}
	projectID := getStringValue(packetProjectIDVar, *projectFlag, "")
	if projectID == "" {
		log.Fatal("Must specify a Packet.net Project ID with --project-id")
	}
	hostname := getStringValue(packetHostnameVar, *hostNameFlag, "")
	name := getStringValue(packetNameVar, *nameFlag, prefix)
	osType := "custom_ipxe"
	billing := "hourly"

	if !*keepFlag && !*consoleFlag {
		log.Fatalf("Combination of keep=%t and console=%t makes little sense", *keepFlag, *consoleFlag)
	}

	// Read kernel command line
	var cmdline string
	if c, err := ioutil.ReadFile(prefix + "-cmdline"); err != nil {
		log.Fatalf("Cannot open cmdline file: %v", err)
	} else {
		cmdline = string(c)
	}

	// Serve files with a local http server
	var httpServer *http.Server
	if *serveFlag != "" {
		fs := serveFiles{[]string{fmt.Sprintf("%s-kernel", name), fmt.Sprintf("%s-initrd.img", name)}}
		httpServer = &http.Server{Addr: ":8080", Handler: http.FileServer(fs)}
		go func() {
			log.Debugf("Listening on http://%s\n", *serveFlag)
			if err := httpServer.ListenAndServe(); err != nil {
				log.Infof("http server exited with: %v", err)
			}
		}()
	}

	// Build the iPXE script
	// Note, we *append* the <prefix>-cmdline. iXPE booting will
	// need the first set of "kernel-params" and we don't want to
	// require these to be added to every YAML file.
	userData := "#!ipxe\n\n"
	userData += "dhcp\n"
	userData += fmt.Sprintf("set base-url %s\n", url)
	userData += fmt.Sprintf("set kernel-params ip=dhcp nomodeset ro serial console=ttyS1,115200 %s\n", cmdline)
	userData += fmt.Sprintf("kernel ${base-url}/%s-kernel ${kernel-params}\n", name)
	userData += fmt.Sprintf("initrd ${base-url}/%s-initrd.img\n", name)
	userData += "boot"
	log.Debugf("Using userData of:\n%s\n", userData)

	// Make sure the URL works
	initrdURL := fmt.Sprintf("%s/%s-initrd.img", url, name)
	kernelURL := fmt.Sprintf("%s/%s-kernel", url, name)
	validateHTTPURL(kernelURL)
	validateHTTPURL(initrdURL)

	client := packngo.NewClient("", apiKey, nil)
	tags := []string{}
	req := packngo.DeviceCreateRequest{
		HostName:     hostname,
		Plan:         plan,
		Facility:     facility,
		OS:           osType,
		BillingCycle: billing,
		ProjectID:    projectID,
		UserData:     userData,
		Tags:         tags,
		AlwaysPXE:    *alwaysPXE,
	}
	dev, _, err := client.Devices.Create(&req)
	if err != nil {
		log.Fatal(err)
	}
	b, err := json.MarshalIndent(dev, "", "    ")
	if err != nil {
		log.Fatal(err)
	}
	log.Debugf("%s\n", string(b))

	log.Printf("Machine booting...")

	sshHost := "sos." + dev.Facility.Code + ".packet.net"
	if *consoleFlag {
		// Connect to the serial console
		if err := sshSOS(dev.ID, sshHost); err != nil {
			log.Fatal(err)
		}
	} else {
		log.Printf("Access the console with: ssh %s@%s", dev.ID, sshHost)

		// if the serve option is present, wait till 'ctrl-c' is hit.
		// Otherwise we wouldn't serve the files
		if *serveFlag != "" {
			stop := make(chan os.Signal, 1)
			signal.Notify(stop, os.Interrupt)
			log.Printf("Hit ctrl-c to stop http server")
			<-stop
		}
	}

	// Stop the http server before exiting
	if *serveFlag != "" {
		log.Debugf("Shutting down http server...")
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		httpServer.Shutdown(ctx)
	}

	if !*keepFlag {
		if _, err := client.Devices.Delete(dev.ID); err != nil {
			log.Fatalf("Unable to delete device: %v", err)
		}
	}

}

// validateHTTPURL does a sanity check that a URL returns a 200 or 300 response
func validateHTTPURL(url string) {
	log.Infof("Validating URL: %s", url)
	resp, err := http.Head(url)
	if err != nil {
		log.Fatal(err)
	}
	if resp.StatusCode >= 400 {
		log.Fatal("Got a non 200- or 300- HTTP response code: %s", resp)
	}
	log.Debugf("OK: %d response code", resp.StatusCode)
}

func sshSOS(user, host string) error {
	log.Debugf("console: ssh %s@%s", user, host)

	hostKey, err := sshHostKey(host)
	if err != nil {
		return fmt.Errorf("Host key not found. Maybe need to add it? %v", err)
	}

	sshConfig := &ssh.ClientConfig{
		User:            user,
		HostKeyCallback: ssh.FixedHostKey(hostKey),
		Auth: []ssh.AuthMethod{
			sshAgent(),
		},
	}

	c, err := ssh.Dial("tcp", host+":22", sshConfig)
	if err != nil {
		return fmt.Errorf("Failed to dial: %s", err)
	}

	s, err := c.NewSession()
	if err != nil {
		return fmt.Errorf("Failed to create session: %v", err)
	}
	defer s.Close()

	s.Stdout = os.Stdout
	s.Stderr = os.Stderr
	s.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO:  0,
		ssh.IGNCR: 1,
	}

	width, height, err := terminal.GetSize(0)
	if err != nil {
		log.Warningf("Error getting terminal size. Ignored. %v", err)
		width = 80
		height = 40
	}
	if err := s.RequestPty("vt100", width, height, modes); err != nil {
		return fmt.Errorf("Request for PTY failed: %v", err)
	}
	oldState, err := terminal.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return err
	}
	defer terminal.Restore(0, oldState)

	// Start remote shell
	if err := s.Shell(); err != nil {
		return fmt.Errorf("Failed to start shell: %v", err)
	}

	s.Wait()
	return nil
}

// Get a ssh-agent AuthMethod
func sshAgent() ssh.AuthMethod {
	sshAgent, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		log.Fatalf("Failed to dial ssh-agent: %v", err)
	}
	return ssh.PublicKeysCallback(agent.NewClient(sshAgent).Signers)
}

// This function returns the host key for a given host (the SOS server).
// If it can't be found, it errors
func sshHostKey(host string) (ssh.PublicKey, error) {
	f, err := os.Open(filepath.Join(os.Getenv("HOME"), ".ssh", "known_hosts"))
	if err != nil {
		return nil, fmt.Errorf("Can't open know_hosts file: %v", err)
	}
	defer f.Close()

	s := bufio.NewScanner(f)

	var hostKey ssh.PublicKey
	for s.Scan() {
		fields := strings.Split(s.Text(), " ")
		if len(fields) != 3 {
			continue
		}
		if strings.Contains(fields[0], host) {
			var err error
			hostKey, _, _, _, err = ssh.ParseAuthorizedKey(s.Bytes())
			if err != nil {
				return nil, fmt.Errorf("Error parsing %q: %v", fields[2], err)
			}
			break
		}
	}

	if hostKey == nil {
		return nil, fmt.Errorf("No hostkey for %s", host)
	}
	return hostKey, nil
}

// This implements a http.FileSystem which only responds to specific files.
type serveFiles struct {
	files []string
}

// Open implements the Open method for the serveFiles FileSystem
// implementation.
// It converts both the name from the URL and the files provided in
// the serveFiles structure into cleaned, absolute filesystem path and
// only returns the file if the requested name matches one of the
// files in the list.
func (fs serveFiles) Open(name string) (http.File, error) {
	cwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	name = filepath.Join(cwd, filepath.FromSlash(path.Clean("/"+name)))
	for _, fn := range fs.files {
		fn = filepath.Join(cwd, filepath.FromSlash(path.Clean("/"+fn)))
		if name == fn {
			f, err := os.Open(fn)
			if err != nil {
				return nil, err
			}
			log.Debugf("Serving: %s", fn)
			return f, nil
		}
	}
	return nil, fmt.Errorf("File %s not found", name)
}
<|endoftext|>


<commit_before><commit_msg>Handle zero values<commit_after><|endoftext|>


<commit_before>package providers

import "github.com/drgarcia1986/gonews/story"

type StoryRequest struct {
	Story *story.Story
	Err   error
}

type Provider interface {
	GetStories(int, int) (<-chan chan *StoryRequest, error)
}
<commit_msg>add method `Name()` to providers interface<commit_after>package providers

import "github.com/drgarcia1986/gonews/story"

type StoryRequest struct {
	Story *story.Story
	Err   error
}

type Provider interface {
	Name() string
	GetStories(int, int) (<-chan chan *StoryRequest, error)
}
<|endoftext|>


<commit_before>/*
Copyright 2021 Gravitational, 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 clusters

import (
	"context"

	"github.com/gravitational/teleport/api/client/proto"
	"github.com/gravitational/teleport/api/defaults"
	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/lib/auth"
	"github.com/gravitational/teleport/lib/client"
	dbprofile "github.com/gravitational/teleport/lib/client/db"
	libdefaults "github.com/gravitational/teleport/lib/defaults"
	"github.com/gravitational/teleport/lib/services"
	"github.com/gravitational/teleport/lib/teleterm/api/uri"
	"github.com/gravitational/teleport/lib/tlsca"

	"github.com/gravitational/trace"
)

// Database describes database
type Database struct {
	// URI is the database URI
	URI uri.ResourceURI
	types.Database
}

// GetDatabase returns a database
func (c *Cluster) GetDatabase(ctx context.Context, dbURI string) (*Database, error) {
	dbs, err := c.GetDatabases(ctx)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	for _, db := range dbs {
		if db.URI.String() == dbURI {
			return &db, nil
		}
	}

	return nil, trace.NotFound("database is not found: %v", dbURI)
}

// GetDatabases returns databases
func (c *Cluster) GetDatabases(ctx context.Context) ([]Database, error) {
	var dbs []types.Database
	err := addMetadataToRetryableError(ctx, func() error {
		proxyClient, err := c.clusterClient.ConnectToProxy(ctx)
		if err != nil {
			return trace.Wrap(err)
		}
		defer proxyClient.Close()

		dbs, err = proxyClient.FindDatabasesByFilters(ctx, proto.ListResourcesRequest{
			Namespace: defaults.Namespace,
		})
		if err != nil {
			return trace.Wrap(err)
		}

		return nil
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}

	var responseDbs []Database
	for _, db := range dbs {
		responseDbs = append(responseDbs, Database{
			URI:      c.URI.AppendDB(db.GetName()),
			Database: db,
		})
	}

	return responseDbs, nil
}

// ReissueDBCerts issues new certificates for specific DB access
func (c *Cluster) ReissueDBCerts(ctx context.Context, user string, db types.Database) error {
	// When generating certificate for MongoDB access, database username must
	// be encoded into it. This is required to be able to tell which database
	// user to authenticate the connection as.
	if db.GetProtocol() == libdefaults.ProtocolMongoDB && user == "" {
		return trace.BadParameter("please provide the database user name using --db-user flag")
	}

	err := addMetadataToRetryableError(ctx, func() error {
		// Refresh the certs to account for clusterClient.SiteName pointing at a leaf cluster.
		err := c.clusterClient.ReissueUserCerts(ctx, client.CertCacheKeep, client.ReissueParams{
			RouteToCluster: c.clusterClient.SiteName,
			AccessRequests: c.status.ActiveRequests.AccessRequests,
		})
		if err != nil {
			return trace.Wrap(err)
		}

		// Fetch the certs for the database.
		err = c.clusterClient.ReissueUserCerts(ctx, client.CertCacheKeep, client.ReissueParams{
			RouteToCluster: c.clusterClient.SiteName,
			RouteToDatabase: proto.RouteToDatabase{
				ServiceName: db.GetName(),
				Protocol:    db.GetProtocol(),
				Username:    user,
			},
			AccessRequests: c.status.ActiveRequests.AccessRequests,
		})
		if err != nil {
			return trace.Wrap(err)
		}

		return nil
	})
	if err != nil {
		return trace.Wrap(err)
	}

	// Update the database-specific connection profile file.
	err = dbprofile.Add(ctx, c.clusterClient, tlsca.RouteToDatabase{
		ServiceName: db.GetName(),
		Protocol:    db.GetProtocol(),
		Username:    user,
	}, c.status)
	if err != nil {
		return trace.Wrap(err)
	}

	return nil
}

// GetAllowedDatabaseUsers returns allowed users for the given database based on the role set.
func (c *Cluster) GetAllowedDatabaseUsers(ctx context.Context, dbURI string) ([]string, error) {
	var authClient auth.ClientI

	err := addMetadataToRetryableError(ctx, func() error {
		proxyClient, err := c.clusterClient.ConnectToProxy(ctx)
		if err != nil {
			return trace.Wrap(err)
		}
		defer proxyClient.Close()

		authClient, err = proxyClient.ConnectToCluster(ctx, c.clusterClient.SiteName)
		if err != nil {
			return trace.Wrap(err)
		}

		return nil
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}

	defer authClient.Close()

	roleSet, err := services.FetchAllClusterRoles(ctx, authClient, c.status.Roles, c.status.Traits)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	db, err := c.GetDatabase(ctx, dbURI)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	dbUsers := roleSet.EnumerateDatabaseUsers(db)

	return dbUsers.Allowed(), nil
}
<commit_msg>Connect: Fix premature proxyClient.Close() when getting db users (#14230)<commit_after>/*
Copyright 2021 Gravitational, 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 clusters

import (
	"context"

	"github.com/gravitational/teleport/api/client/proto"
	"github.com/gravitational/teleport/api/defaults"
	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/lib/auth"
	"github.com/gravitational/teleport/lib/client"
	dbprofile "github.com/gravitational/teleport/lib/client/db"
	libdefaults "github.com/gravitational/teleport/lib/defaults"
	"github.com/gravitational/teleport/lib/services"
	"github.com/gravitational/teleport/lib/teleterm/api/uri"
	"github.com/gravitational/teleport/lib/tlsca"

	"github.com/gravitational/trace"
)

// Database describes database
type Database struct {
	// URI is the database URI
	URI uri.ResourceURI
	types.Database
}

// GetDatabase returns a database
func (c *Cluster) GetDatabase(ctx context.Context, dbURI string) (*Database, error) {
	dbs, err := c.GetDatabases(ctx)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	for _, db := range dbs {
		if db.URI.String() == dbURI {
			return &db, nil
		}
	}

	return nil, trace.NotFound("database is not found: %v", dbURI)
}

// GetDatabases returns databases
func (c *Cluster) GetDatabases(ctx context.Context) ([]Database, error) {
	var dbs []types.Database
	err := addMetadataToRetryableError(ctx, func() error {
		proxyClient, err := c.clusterClient.ConnectToProxy(ctx)
		if err != nil {
			return trace.Wrap(err)
		}
		defer proxyClient.Close()

		dbs, err = proxyClient.FindDatabasesByFilters(ctx, proto.ListResourcesRequest{
			Namespace: defaults.Namespace,
		})
		if err != nil {
			return trace.Wrap(err)
		}

		return nil
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}

	var responseDbs []Database
	for _, db := range dbs {
		responseDbs = append(responseDbs, Database{
			URI:      c.URI.AppendDB(db.GetName()),
			Database: db,
		})
	}

	return responseDbs, nil
}

// ReissueDBCerts issues new certificates for specific DB access
func (c *Cluster) ReissueDBCerts(ctx context.Context, user string, db types.Database) error {
	// When generating certificate for MongoDB access, database username must
	// be encoded into it. This is required to be able to tell which database
	// user to authenticate the connection as.
	if db.GetProtocol() == libdefaults.ProtocolMongoDB && user == "" {
		return trace.BadParameter("please provide the database user name using --db-user flag")
	}

	err := addMetadataToRetryableError(ctx, func() error {
		// Refresh the certs to account for clusterClient.SiteName pointing at a leaf cluster.
		err := c.clusterClient.ReissueUserCerts(ctx, client.CertCacheKeep, client.ReissueParams{
			RouteToCluster: c.clusterClient.SiteName,
			AccessRequests: c.status.ActiveRequests.AccessRequests,
		})
		if err != nil {
			return trace.Wrap(err)
		}

		// Fetch the certs for the database.
		err = c.clusterClient.ReissueUserCerts(ctx, client.CertCacheKeep, client.ReissueParams{
			RouteToCluster: c.clusterClient.SiteName,
			RouteToDatabase: proto.RouteToDatabase{
				ServiceName: db.GetName(),
				Protocol:    db.GetProtocol(),
				Username:    user,
			},
			AccessRequests: c.status.ActiveRequests.AccessRequests,
		})
		if err != nil {
			return trace.Wrap(err)
		}

		return nil
	})
	if err != nil {
		return trace.Wrap(err)
	}

	// Update the database-specific connection profile file.
	err = dbprofile.Add(ctx, c.clusterClient, tlsca.RouteToDatabase{
		ServiceName: db.GetName(),
		Protocol:    db.GetProtocol(),
		Username:    user,
	}, c.status)
	if err != nil {
		return trace.Wrap(err)
	}

	return nil
}

// GetAllowedDatabaseUsers returns allowed users for the given database based on the role set.
func (c *Cluster) GetAllowedDatabaseUsers(ctx context.Context, dbURI string) ([]string, error) {
	var authClient auth.ClientI
	var proxyClient *client.ProxyClient
	var err error

	err = addMetadataToRetryableError(ctx, func() error {
		proxyClient, err = c.clusterClient.ConnectToProxy(ctx)
		if err != nil {
			return trace.Wrap(err)
		}

		return nil
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer proxyClient.Close()

	authClient, err = proxyClient.ConnectToCluster(ctx, c.clusterClient.SiteName)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer authClient.Close()

	roleSet, err := services.FetchAllClusterRoles(ctx, authClient, c.status.Roles, c.status.Traits)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	db, err := c.GetDatabase(ctx, dbURI)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	dbUsers := roleSet.EnumerateDatabaseUsers(db)

	return dbUsers.Allowed(), nil
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2018 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

// Package httpmultibin is indended only for use in tests, do not import in production code!
package httpmultibin

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/andybalholm/brotli"
	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
	"github.com/mccutchen/go-httpbin/httpbin"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	grpctest "google.golang.org/grpc/test/grpc_testing"

	"go.k6.io/k6/lib"
	"go.k6.io/k6/lib/netext"
	"go.k6.io/k6/lib/netext/httpext"
	"go.k6.io/k6/lib/types"
)

// GetTLSClientConfig returns a TLS config that trusts the supplied
// httptest.Server certificate as well as all the system root certificates
func GetTLSClientConfig(t testing.TB, srv *httptest.Server) *tls.Config {
	var err error

	certs := x509.NewCertPool()

	if runtime.GOOS != "windows" {
		certs, err = x509.SystemCertPool()
		require.NoError(t, err)
	}

	for _, c := range srv.TLS.Certificates {
		roots, err := x509.ParseCertificates(c.Certificate[len(c.Certificate)-1])
		require.NoError(t, err)
		for _, root := range roots {
			certs.AddCert(root)
		}
	}
	return &tls.Config{
		RootCAs:            certs,
		InsecureSkipVerify: false,
		Renegotiation:      tls.RenegotiateFreelyAsClient,
	}
}

const httpDomain = "httpbin.local"

// We have to use example.com if we want a real HTTPS domain with a valid
// certificate because the default httptest certificate is for example.com:
// https://golang.org/src/net/http/internal/testcert.go?s=399:410#L10
const httpsDomain = "example.com"

// HTTPMultiBin can be used as a local alternative of httpbin.org. It offers both http and https servers, as well as real domains
type HTTPMultiBin struct {
	Mux             *http.ServeMux
	ServerHTTP      *httptest.Server
	ServerHTTPS     *httptest.Server
	ServerHTTP2     *httptest.Server
	ServerGRPC      *grpc.Server
	GRPCStub        *GRPCStub
	Replacer        *strings.Replacer
	TLSClientConfig *tls.Config
	Dialer          *netext.Dialer
	HTTPTransport   *http.Transport
	Context         context.Context
}

type jsonBody struct {
	Header      http.Header `json:"headers"`
	Compression string      `json:"compression"`
}

func getWebsocketHandler(echo bool, closePrematurely bool) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
		if err != nil {
			return
		}
		if echo {
			messageType, r, e := conn.NextReader()
			if e != nil {
				return
			}
			var wc io.WriteCloser
			wc, err = conn.NextWriter(messageType)
			if err != nil {
				return
			}
			if _, err = io.Copy(wc, r); err != nil {
				return
			}
			if err = wc.Close(); err != nil {
				return
			}
		}
		// closePrematurely=true mimics an invalid WS server that doesn't
		// send a close control frame before closing the connection.
		if !closePrematurely {
			closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
			_ = conn.WriteControl(websocket.CloseMessage, closeMsg, time.Now().Add(time.Second))
			// Wait for response control frame
			<-time.After(time.Second)
		}
		err = conn.Close()
		if err != nil {
			return
		}
	})
}

func writeJSON(w io.Writer, v interface{}) error {
	e := json.NewEncoder(w)
	e.SetIndent("", "  ")
	if err := e.Encode(v); err != nil {
		return fmt.Errorf("failed to encode JSON: %w", err)
	}
	return nil
}

func getEncodedHandler(t testing.TB, compressionType httpext.CompressionType) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		var (
			encoding string
			err      error
			encw     io.WriteCloser
		)

		switch compressionType {
		case httpext.CompressionTypeBr:
			encw = brotli.NewWriter(rw)
			encoding = "br"
		case httpext.CompressionTypeZstd:
			encw, _ = zstd.NewWriter(rw)
			encoding = "zstd"
		}

		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}
		err = writeJSON(encw, data)
		if encw != nil {
			_ = encw.Close()
		}
		require.NoError(t, err)
	})
}

func getZstdBrHandler(t testing.TB) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		encoding := "zstd, br"
		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}

		bw := brotli.NewWriter(rw)
		zw, _ := zstd.NewWriter(bw)
		defer func() {
			_ = zw.Close()
			_ = bw.Close()
		}()

		require.NoError(t, writeJSON(zw, data))
	})
}

// GRPCStub is an easily customisable TestServiceServer
type GRPCStub struct {
	grpctest.TestServiceServer
	EmptyCallFunc func(context.Context, *grpctest.Empty) (*grpctest.Empty, error)
	UnaryCallFunc func(context.Context, *grpctest.SimpleRequest) (*grpctest.SimpleResponse, error)
}

// EmptyCall implements the interface for the gRPC TestServiceServer
func (s *GRPCStub) EmptyCall(ctx context.Context, req *grpctest.Empty) (*grpctest.Empty, error) {
	if s.EmptyCallFunc != nil {
		return s.EmptyCallFunc(ctx, req)
	}

	return nil, status.Errorf(codes.Unimplemented, "method EmptyCall not implemented")
}

// UnaryCall implements the interface for the gRPC TestServiceServer
func (s *GRPCStub) UnaryCall(ctx context.Context, req *grpctest.SimpleRequest) (*grpctest.SimpleResponse, error) {
	if s.UnaryCallFunc != nil {
		return s.UnaryCallFunc(ctx, req)
	}

	return nil, status.Errorf(codes.Unimplemented, "method UnaryCall not implemented")
}

// StreamingOutputCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) StreamingOutputCall(*grpctest.StreamingOutputCallRequest,
	grpctest.TestService_StreamingOutputCallServer) error {
	return status.Errorf(codes.Unimplemented, "method StreamingOutputCall not implemented")
}

// StreamingInputCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) StreamingInputCall(grpctest.TestService_StreamingInputCallServer) error {
	return status.Errorf(codes.Unimplemented, "method StreamingInputCall not implemented")
}

// FullDuplexCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) FullDuplexCall(grpctest.TestService_FullDuplexCallServer) error {
	return status.Errorf(codes.Unimplemented, "method FullDuplexCall not implemented")
}

// HalfDuplexCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) HalfDuplexCall(grpctest.TestService_HalfDuplexCallServer) error {
	return status.Errorf(codes.Unimplemented, "method HalfDuplexCall not implemented")
}

// NewHTTPMultiBin returns a fully configured and running HTTPMultiBin
func NewHTTPMultiBin(t testing.TB) *HTTPMultiBin {
	// Create a http.ServeMux and set the httpbin handler as the default
	mux := http.NewServeMux()
	mux.Handle("/brotli", getEncodedHandler(t, httpext.CompressionTypeBr))
	mux.Handle("/ws-echo", getWebsocketHandler(true, false))
	mux.Handle("/ws-echo-invalid", getWebsocketHandler(true, true))
	mux.Handle("/ws-close", getWebsocketHandler(false, false))
	mux.Handle("/ws-close-invalid", getWebsocketHandler(false, true))
	mux.Handle("/zstd", getEncodedHandler(t, httpext.CompressionTypeZstd))
	mux.Handle("/zstd-br", getZstdBrHandler(t))
	mux.Handle("/", httpbin.New().Handler())

	// Initialize the HTTP server and get its details
	httpSrv := httptest.NewServer(mux)
	httpURL, err := url.Parse(httpSrv.URL)
	require.NoError(t, err)
	httpIP := net.ParseIP(httpURL.Hostname())
	require.NotNil(t, httpIP)

	// Initialize the HTTPS server and get its details and tls config
	httpsSrv := httptest.NewTLSServer(mux)
	httpsURL, err := url.Parse(httpsSrv.URL)
	require.NoError(t, err)
	httpsIP := net.ParseIP(httpsURL.Hostname())
	require.NotNil(t, httpsIP)
	tlsConfig := GetTLSClientConfig(t, httpsSrv)

	// Initialize the gRPC server
	grpcSrv := grpc.NewServer()
	stub := &GRPCStub{}
	grpctest.RegisterTestServiceServer(grpcSrv, stub)

	cmux := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcSrv.ServeHTTP(w, r)

			return
		}
		mux.ServeHTTP(w, r)
	})

	// Initialize the HTTP2 server, with a copy of the https tls config
	http2Srv := httptest.NewUnstartedServer(cmux)
	http2Srv.EnableHTTP2 = true
	http2Srv.TLS = &(*tlsConfig) // copy it
	http2Srv.TLS.NextProtos = []string{http2.NextProtoTLS}
	require.NoError(t, err)
	http2Srv.StartTLS()
	http2URL, err := url.Parse(http2Srv.URL)
	require.NoError(t, err)
	http2IP := net.ParseIP(http2URL.Hostname())
	require.NotNil(t, http2IP)

	httpDomainValue, err := lib.NewHostAddress(httpIP, "")
	require.NoError(t, err)
	httpsDomainValue, err := lib.NewHostAddress(httpsIP, "")
	require.NoError(t, err)

	// Set up the dialer with shorter timeouts and the custom domains
	dialer := netext.NewDialer(net.Dialer{
		Timeout:   2 * time.Second,
		KeepAlive: 10 * time.Second,
		DualStack: true,
	}, netext.NewResolver(net.LookupIP, 0, types.DNSfirst, types.DNSpreferIPv4))
	dialer.Hosts = map[string]*lib.HostAddress{
		httpDomain:  httpDomainValue,
		httpsDomain: httpsDomainValue,
	}

	// Pre-configure the HTTP client transport with the dialer and TLS config (incl. HTTP2 support)
	transport := &http.Transport{
		DialContext:     dialer.DialContext,
		TLSClientConfig: tlsConfig,
	}
	require.NoError(t, http2.ConfigureTransport(transport))

	ctx, ctxCancel := context.WithCancel(context.Background())

	result := &HTTPMultiBin{
		Mux:         mux,
		ServerHTTP:  httpSrv,
		ServerHTTPS: httpsSrv,
		ServerHTTP2: http2Srv,
		ServerGRPC:  grpcSrv,
		GRPCStub:    stub,
		Replacer: strings.NewReplacer(
			"HTTPBIN_IP_URL", httpSrv.URL,
			"HTTPBIN_DOMAIN", httpDomain,
			"HTTPBIN_URL", fmt.Sprintf("http://%s:%s", httpDomain, httpURL.Port()),
			"WSBIN_URL", fmt.Sprintf("ws://%s:%s", httpDomain, httpURL.Port()),
			"HTTPBIN_IP", httpIP.String(),
			"HTTPBIN_PORT", httpURL.Port(),
			"HTTPSBIN_IP_URL", httpsSrv.URL,
			"HTTPSBIN_DOMAIN", httpsDomain,
			"HTTPSBIN_URL", fmt.Sprintf("https://%s:%s", httpsDomain, httpsURL.Port()),
			"WSSBIN_URL", fmt.Sprintf("wss://%s:%s", httpsDomain, httpsURL.Port()),
			"HTTPSBIN_IP", httpsIP.String(),
			"HTTPSBIN_PORT", httpsURL.Port(),

			"HTTP2BIN_IP_URL", http2Srv.URL,
			"HTTP2BIN_DOMAIN", httpsDomain,
			"HTTP2BIN_URL", fmt.Sprintf("https://%s:%s", httpsDomain, http2URL.Port()),
			"HTTP2BIN_IP", http2IP.String(),
			"HTTP2BIN_PORT", http2URL.Port(),

			"GRPCBIN_ADDR", fmt.Sprintf("%s:%s", httpsDomain, http2URL.Port()),
		),
		TLSClientConfig: tlsConfig,
		Dialer:          dialer,
		HTTPTransport:   transport,
		Context:         ctx,
	}

	t.Cleanup(func() {
		grpcSrv.Stop()
		http2Srv.Close()
		httpsSrv.Close()
		httpSrv.Close()
		ctxCancel()
	})
	return result
}
<commit_msg>Drop httpmultibin dependency on httpext<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2018 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

// Package httpmultibin is indended only for use in tests, do not import in production code!
package httpmultibin

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/andybalholm/brotli"
	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
	"github.com/mccutchen/go-httpbin/httpbin"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	grpctest "google.golang.org/grpc/test/grpc_testing"

	"go.k6.io/k6/lib"
	"go.k6.io/k6/lib/netext"
	"go.k6.io/k6/lib/types"
)

// GetTLSClientConfig returns a TLS config that trusts the supplied
// httptest.Server certificate as well as all the system root certificates
func GetTLSClientConfig(t testing.TB, srv *httptest.Server) *tls.Config {
	var err error

	certs := x509.NewCertPool()

	if runtime.GOOS != "windows" {
		certs, err = x509.SystemCertPool()
		require.NoError(t, err)
	}

	for _, c := range srv.TLS.Certificates {
		roots, err := x509.ParseCertificates(c.Certificate[len(c.Certificate)-1])
		require.NoError(t, err)
		for _, root := range roots {
			certs.AddCert(root)
		}
	}
	return &tls.Config{
		RootCAs:            certs,
		InsecureSkipVerify: false,
		Renegotiation:      tls.RenegotiateFreelyAsClient,
	}
}

const httpDomain = "httpbin.local"

// We have to use example.com if we want a real HTTPS domain with a valid
// certificate because the default httptest certificate is for example.com:
// https://golang.org/src/net/http/internal/testcert.go?s=399:410#L10
const httpsDomain = "example.com"

// HTTPMultiBin can be used as a local alternative of httpbin.org. It offers both http and https servers, as well as real domains
type HTTPMultiBin struct {
	Mux             *http.ServeMux
	ServerHTTP      *httptest.Server
	ServerHTTPS     *httptest.Server
	ServerHTTP2     *httptest.Server
	ServerGRPC      *grpc.Server
	GRPCStub        *GRPCStub
	Replacer        *strings.Replacer
	TLSClientConfig *tls.Config
	Dialer          *netext.Dialer
	HTTPTransport   *http.Transport
	Context         context.Context
}

type jsonBody struct {
	Header      http.Header `json:"headers"`
	Compression string      `json:"compression"`
}

func getWebsocketHandler(echo bool, closePrematurely bool) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, req, w.Header())
		if err != nil {
			return
		}
		if echo {
			messageType, r, e := conn.NextReader()
			if e != nil {
				return
			}
			var wc io.WriteCloser
			wc, err = conn.NextWriter(messageType)
			if err != nil {
				return
			}
			if _, err = io.Copy(wc, r); err != nil {
				return
			}
			if err = wc.Close(); err != nil {
				return
			}
		}
		// closePrematurely=true mimics an invalid WS server that doesn't
		// send a close control frame before closing the connection.
		if !closePrematurely {
			closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
			_ = conn.WriteControl(websocket.CloseMessage, closeMsg, time.Now().Add(time.Second))
			// Wait for response control frame
			<-time.After(time.Second)
		}
		err = conn.Close()
		if err != nil {
			return
		}
	})
}

func writeJSON(w io.Writer, v interface{}) error {
	e := json.NewEncoder(w)
	e.SetIndent("", "  ")
	if err := e.Encode(v); err != nil {
		return fmt.Errorf("failed to encode JSON: %w", err)
	}
	return nil
}

func getEncodedHandler(t testing.TB, compressionType string) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		var (
			encoding string
			err      error
			encw     io.WriteCloser
		)

		switch compressionType {
		case "br":
			encw = brotli.NewWriter(rw)
			encoding = "br"
		case "zstd":
			encw, _ = zstd.NewWriter(rw)
			encoding = "zstd"
		}

		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}
		err = writeJSON(encw, data)
		if encw != nil {
			_ = encw.Close()
		}
		require.NoError(t, err)
	})
}

func getZstdBrHandler(t testing.TB) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		encoding := "zstd, br"
		rw.Header().Set("Content-Type", "application/json")
		rw.Header().Add("Content-Encoding", encoding)
		data := jsonBody{
			Header:      req.Header,
			Compression: encoding,
		}

		bw := brotli.NewWriter(rw)
		zw, _ := zstd.NewWriter(bw)
		defer func() {
			_ = zw.Close()
			_ = bw.Close()
		}()

		require.NoError(t, writeJSON(zw, data))
	})
}

// GRPCStub is an easily customisable TestServiceServer
type GRPCStub struct {
	grpctest.TestServiceServer
	EmptyCallFunc func(context.Context, *grpctest.Empty) (*grpctest.Empty, error)
	UnaryCallFunc func(context.Context, *grpctest.SimpleRequest) (*grpctest.SimpleResponse, error)
}

// EmptyCall implements the interface for the gRPC TestServiceServer
func (s *GRPCStub) EmptyCall(ctx context.Context, req *grpctest.Empty) (*grpctest.Empty, error) {
	if s.EmptyCallFunc != nil {
		return s.EmptyCallFunc(ctx, req)
	}

	return nil, status.Errorf(codes.Unimplemented, "method EmptyCall not implemented")
}

// UnaryCall implements the interface for the gRPC TestServiceServer
func (s *GRPCStub) UnaryCall(ctx context.Context, req *grpctest.SimpleRequest) (*grpctest.SimpleResponse, error) {
	if s.UnaryCallFunc != nil {
		return s.UnaryCallFunc(ctx, req)
	}

	return nil, status.Errorf(codes.Unimplemented, "method UnaryCall not implemented")
}

// StreamingOutputCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) StreamingOutputCall(*grpctest.StreamingOutputCallRequest,
	grpctest.TestService_StreamingOutputCallServer) error {
	return status.Errorf(codes.Unimplemented, "method StreamingOutputCall not implemented")
}

// StreamingInputCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) StreamingInputCall(grpctest.TestService_StreamingInputCallServer) error {
	return status.Errorf(codes.Unimplemented, "method StreamingInputCall not implemented")
}

// FullDuplexCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) FullDuplexCall(grpctest.TestService_FullDuplexCallServer) error {
	return status.Errorf(codes.Unimplemented, "method FullDuplexCall not implemented")
}

// HalfDuplexCall implements the interface for the gRPC TestServiceServer
func (*GRPCStub) HalfDuplexCall(grpctest.TestService_HalfDuplexCallServer) error {
	return status.Errorf(codes.Unimplemented, "method HalfDuplexCall not implemented")
}

// NewHTTPMultiBin returns a fully configured and running HTTPMultiBin
//nolint:funlen
func NewHTTPMultiBin(t testing.TB) *HTTPMultiBin {
	// Create a http.ServeMux and set the httpbin handler as the default
	mux := http.NewServeMux()
	mux.Handle("/brotli", getEncodedHandler(t, "br"))
	mux.Handle("/ws-echo", getWebsocketHandler(true, false))
	mux.Handle("/ws-echo-invalid", getWebsocketHandler(true, true))
	mux.Handle("/ws-close", getWebsocketHandler(false, false))
	mux.Handle("/ws-close-invalid", getWebsocketHandler(false, true))
	mux.Handle("/zstd", getEncodedHandler(t, "zstd"))
	mux.Handle("/zstd-br", getZstdBrHandler(t))
	mux.Handle("/", httpbin.New().Handler())

	// Initialize the HTTP server and get its details
	httpSrv := httptest.NewServer(mux)
	httpURL, err := url.Parse(httpSrv.URL)
	require.NoError(t, err)
	httpIP := net.ParseIP(httpURL.Hostname())
	require.NotNil(t, httpIP)

	// Initialize the HTTPS server and get its details and tls config
	httpsSrv := httptest.NewTLSServer(mux)
	httpsURL, err := url.Parse(httpsSrv.URL)
	require.NoError(t, err)
	httpsIP := net.ParseIP(httpsURL.Hostname())
	require.NotNil(t, httpsIP)
	tlsConfig := GetTLSClientConfig(t, httpsSrv)

	// Initialize the gRPC server
	grpcSrv := grpc.NewServer()
	stub := &GRPCStub{}
	grpctest.RegisterTestServiceServer(grpcSrv, stub)

	cmux := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcSrv.ServeHTTP(w, r)

			return
		}
		mux.ServeHTTP(w, r)
	})

	// Initialize the HTTP2 server, with a copy of the https tls config
	http2Srv := httptest.NewUnstartedServer(cmux)
	http2Srv.EnableHTTP2 = true
	http2Srv.TLS = &(*tlsConfig) // copy it
	http2Srv.TLS.NextProtos = []string{http2.NextProtoTLS}
	require.NoError(t, err)
	http2Srv.StartTLS()
	http2URL, err := url.Parse(http2Srv.URL)
	require.NoError(t, err)
	http2IP := net.ParseIP(http2URL.Hostname())
	require.NotNil(t, http2IP)

	httpDomainValue, err := lib.NewHostAddress(httpIP, "")
	require.NoError(t, err)
	httpsDomainValue, err := lib.NewHostAddress(httpsIP, "")
	require.NoError(t, err)

	// Set up the dialer with shorter timeouts and the custom domains
	dialer := netext.NewDialer(net.Dialer{
		Timeout:   2 * time.Second,
		KeepAlive: 10 * time.Second,
		DualStack: true,
	}, netext.NewResolver(net.LookupIP, 0, types.DNSfirst, types.DNSpreferIPv4))
	dialer.Hosts = map[string]*lib.HostAddress{
		httpDomain:  httpDomainValue,
		httpsDomain: httpsDomainValue,
	}

	// Pre-configure the HTTP client transport with the dialer and TLS config (incl. HTTP2 support)
	transport := &http.Transport{
		DialContext:     dialer.DialContext,
		TLSClientConfig: tlsConfig,
	}
	require.NoError(t, http2.ConfigureTransport(transport))

	ctx, ctxCancel := context.WithCancel(context.Background())

	result := &HTTPMultiBin{
		Mux:         mux,
		ServerHTTP:  httpSrv,
		ServerHTTPS: httpsSrv,
		ServerHTTP2: http2Srv,
		ServerGRPC:  grpcSrv,
		GRPCStub:    stub,
		Replacer: strings.NewReplacer(
			"HTTPBIN_IP_URL", httpSrv.URL,
			"HTTPBIN_DOMAIN", httpDomain,
			"HTTPBIN_URL", fmt.Sprintf("http://%s:%s", httpDomain, httpURL.Port()),
			"WSBIN_URL", fmt.Sprintf("ws://%s:%s", httpDomain, httpURL.Port()),
			"HTTPBIN_IP", httpIP.String(),
			"HTTPBIN_PORT", httpURL.Port(),
			"HTTPSBIN_IP_URL", httpsSrv.URL,
			"HTTPSBIN_DOMAIN", httpsDomain,
			"HTTPSBIN_URL", fmt.Sprintf("https://%s:%s", httpsDomain, httpsURL.Port()),
			"WSSBIN_URL", fmt.Sprintf("wss://%s:%s", httpsDomain, httpsURL.Port()),
			"HTTPSBIN_IP", httpsIP.String(),
			"HTTPSBIN_PORT", httpsURL.Port(),

			"HTTP2BIN_IP_URL", http2Srv.URL,
			"HTTP2BIN_DOMAIN", httpsDomain,
			"HTTP2BIN_URL", fmt.Sprintf("https://%s:%s", httpsDomain, http2URL.Port()),
			"HTTP2BIN_IP", http2IP.String(),
			"HTTP2BIN_PORT", http2URL.Port(),

			"GRPCBIN_ADDR", fmt.Sprintf("%s:%s", httpsDomain, http2URL.Port()),
		),
		TLSClientConfig: tlsConfig,
		Dialer:          dialer,
		HTTPTransport:   transport,
		Context:         ctx,
	}

	t.Cleanup(func() {
		grpcSrv.Stop()
		http2Srv.Close()
		httpsSrv.Close()
		httpSrv.Close()
		ctxCancel()
	})
	return result
}
<|endoftext|>


<commit_before>package shared

const (
	NullableTypeString1TestTmpl =
`package {{.Pkg}}

/*
 * CODE GENERATED AUTOMATICALLY WITH https://github.com/reiver/gogen-optiontype
 * THIS FILE SHOULD NOT BE EDITED BY HAND
 */

import (
	"fmt"
	"math"
	"testing"
)

func TestNullableTypeString(t *testing.T) {
	tests := []struct{
		Datum    NullableType
		Expected string
	}{
		{
			Datum:                                SomeNullable(math.MinInt64),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", math.MinInt64),
		},
		{
			Datum                               : SomeNullable(-5),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -5),
		},
		{
			Datum                               : SomeNullable(-4),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -4),
		},
		{
			Datum                               : SomeNullable(-3),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -3),
		},
		{
			Datum                               : SomeNullable(-2),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -2),
		},
		{
			Datum                               : SomeNullable(-1),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -1),
		},
		{
			Datum                               : SomeNullable(0),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 0),
		},
		{
			Datum                               : SomeNullable(1),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 1),
		},
		{
			Datum                               : SomeNullable(2),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 2),
		},
		{
			Datum                               : SomeNullable(3),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 3),
		},
		{
			Datum                               : SomeNullable(4),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 4),
		},
		{
			Datum                               : SomeNullable(5),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 5),
		},
		{
			Datum:                                SomeNullable(math.MaxInt64),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", math.MaxInt64),
		},
	}

	const numRandTests = 40
	for i := 0; i < numRandTests; i++ {
		x := randomness.Int63()

		{
			test := struct {
				Datum NullableType
				Expected string
			}{
				Datum:                                SomeNullable(x),
				Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", x),
			}
			tests = append(tests, test)
		}

		{
			y := -x

			test := struct {
				Datum NullableType
				Expected string
			}{
				Datum:                                SomeNullable(y),
				Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", y),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {
		datum := test.Datum

		if expected, actual := test.Expected, datum.String(); expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v", testNumber, expected, actual)
			continue
		}
	}
}

func TestNullableTypeStringNoneNullable(t *testing.T) {

	datum := NoneNullable()

	if expected, actual := "{{.Pkg}}.NoneNullable()", datum.String(); expected != actual {
		t.Errorf("Expected %q, but actually got %q", expected, actual)
		return
	}
}

func TestNullableTypeStringNull(t *testing.T) {

	datum := None()

	if expected, actual := "{{.Pkg}}.Null()", datum.String(); expected != actual {
		t.Errorf("Expected %q, but actually got %q", expected, actual)
		return
	}
}
`
)
<commit_msg>shared drivers<commit_after>package shared

const (
	NullableTypeString1TestTmpl =
`package {{.Pkg}}

/*
 * CODE GENERATED AUTOMATICALLY WITH https://github.com/reiver/gogen-optiontype
 * THIS FILE SHOULD NOT BE EDITED BY HAND
 */

import (
	"fmt"
	"math"
	"testing"
)

func TestNullableTypeString(t *testing.T) {
	tests := []struct{
		Datum    NullableType
		Expected string
	}{
		{
			Datum:                                SomeNullable(math.MinInt64),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", math.MinInt64),
		},
		{
			Datum                               : SomeNullable(-5),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -5),
		},
		{
			Datum                               : SomeNullable(-4),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -4),
		},
		{
			Datum                               : SomeNullable(-3),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -3),
		},
		{
			Datum                               : SomeNullable(-2),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -2),
		},
		{
			Datum                               : SomeNullable(-1),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", -1),
		},
		{
			Datum                               : SomeNullable(0),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 0),
		},
		{
			Datum                               : SomeNullable(1),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 1),
		},
		{
			Datum                               : SomeNullable(2),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 2),
		},
		{
			Datum                               : SomeNullable(3),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 3),
		},
		{
			Datum                               : SomeNullable(4),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 4),
		},
		{
			Datum                               : SomeNullable(5),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", 5),
		},
		{
			Datum:                                SomeNullable(math.MaxInt64),
			Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", math.MaxInt64),
		},
	}

	const numRandTests = 40
	for i := 0; i < numRandTests; i++ {
		x := randomness.Int63()

		{
			test := struct {
				Datum NullableType
				Expected string
			}{
				Datum:                                SomeNullable(x),
				Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", x),
			}
			tests = append(tests, test)
		}

		{
			y := -x

			test := struct {
				Datum NullableType
				Expected string
			}{
				Datum:                                SomeNullable(y),
				Expected: fmt.Sprintf("{{.Pkg}}.SomeNullable(%d)", y),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {
		datum := test.Datum

		if expected, actual := test.Expected, datum.String(); expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v", testNumber, expected, actual)
			continue
		}
	}
}

func TestNullableTypeStringNoneNullable(t *testing.T) {

	datum := NoneNullable()

	if expected, actual := "{{.Pkg}}.NoneNullable()", datum.String(); expected != actual {
		t.Errorf("Expected %q, but actually got %q", expected, actual)
		return
	}
}

func TestNullableTypeStringNull(t *testing.T) {

	datum := Null()

	if expected, actual := "{{.Pkg}}.Null()", datum.String(); expected != actual {
		t.Errorf("Expected %q, but actually got %q", expected, actual)
		return
	}
}
`
)
<|endoftext|>


<commit_before>// Copyright (c) 2012-2014 Jeremy Latt
// Copyright (c) 2016-2017 Daniel Oaks <daniel@danieloaks.net>
// released under the MIT license

package irc

import (
	"errors"
	"io"
	"strings"
	"sync"

	"github.com/oragono/oragono/irc/utils"
)

var (
	errSendQExceeded = errors.New("SendQ exceeded")

	sendQExceededMessage = []byte("\r\nERROR :SendQ Exceeded\r\n")
)

// Socket represents an IRC socket.
type Socket struct {
	sync.Mutex

	conn IRCConn

	maxSendQBytes int

	// this is a trylock enforcing that only one goroutine can write to `conn` at a time
	writerSemaphore utils.Semaphore

	buffers       [][]byte
	totalLength   int
	closed        bool
	sendQExceeded bool
	finalData     []byte // what to send when we die
	finalized     bool
}

// NewSocket returns a new Socket.
func NewSocket(conn IRCConn, maxSendQBytes int) *Socket {
	result := Socket{
		conn:          conn,
		maxSendQBytes: maxSendQBytes,
	}
	result.writerSemaphore.Initialize(1)
	return &result
}

// Close stops a Socket from being able to send/receive any more data.
func (socket *Socket) Close() {
	socket.Lock()
	socket.closed = true
	socket.Unlock()

	socket.wakeWriter()
}

// Read returns a single IRC line from a Socket.
func (socket *Socket) Read() (string, error) {
	if socket.IsClosed() {
		return "", io.EOF
	}

	lineBytes, err := socket.conn.ReadLine()

	// convert bytes to string
	line := string(lineBytes)

	// read last message properly (such as ERROR/QUIT/etc), just fail next reads/writes
	if err == io.EOF {
		socket.Close()
	}

	if err == io.EOF && strings.TrimSpace(line) != "" {
		// don't do anything
	} else if err != nil {
		return "", err
	}

	return line, nil
}

// Write sends the given string out of Socket. Requirements:
// 1. MUST NOT block for macroscopic amounts of time
// 2. MUST NOT reorder messages
// 3. MUST provide mutual exclusion for socket.conn.Write
// 4. SHOULD NOT tie up additional goroutines, beyond the one blocked on socket.conn.Write
func (socket *Socket) Write(data []byte) (err error) {
	if len(data) == 0 {
		return
	}

	socket.Lock()
	if socket.closed {
		err = io.EOF
	} else {
		prospectiveLen := socket.totalLength + len(data)
		if prospectiveLen > socket.maxSendQBytes {
			socket.sendQExceeded = true
			socket.closed = true
			err = errSendQExceeded
		} else {
			socket.buffers = append(socket.buffers, data)
			socket.totalLength = prospectiveLen
		}
	}
	socket.Unlock()

	socket.wakeWriter()
	return
}

// BlockingWrite sends the given string out of Socket. Requirements:
// 1. MUST block until the message is sent
// 2. MUST bypass sendq (calls to BlockingWrite cannot, on their own, cause a sendq overflow)
// 3. MUST provide mutual exclusion for socket.conn.Write
// 4. MUST respect the same ordering guarantees as Write (i.e., if a call to Write that sends
//    message m1 happens-before a call to BlockingWrite that sends message m2,
//    m1 must be sent on the wire before m2
// Callers MUST be writing to the client's socket from the client's own goroutine;
// other callers must use the nonblocking Write call instead. Otherwise, a client
// with a slow/unreliable connection risks stalling the progress of the system as a whole.
func (socket *Socket) BlockingWrite(data []byte) (err error) {
	if len(data) == 0 {
		return
	}

	// after releasing the semaphore, we must check for fresh data, same as `send`
	defer func() {
		if socket.readyToWrite() {
			socket.wakeWriter()
		}
	}()

	// blocking acquire of the trylock
	socket.writerSemaphore.Acquire()
	defer socket.writerSemaphore.Release()

	// first, flush any buffered data, to preserve the ordering guarantees
	closed := socket.performWrite()
	if closed {
		return io.EOF
	}

	err = socket.conn.Write(data)
	if err != nil {
		socket.finalize()
	}
	return
}

// wakeWriter starts the goroutine that actually performs the write, without blocking
func (socket *Socket) wakeWriter() {
	if socket.writerSemaphore.TryAcquire() {
		// acquired the trylock; send() will release it
		go socket.send()
	}
	// else: do nothing, the holder will check for more data after releasing it
}

// SetFinalData sets the final data to send when the SocketWriter closes.
func (socket *Socket) SetFinalData(data []byte) {
	socket.Lock()
	defer socket.Unlock()
	socket.finalData = data
}

// IsClosed returns whether the socket is closed.
func (socket *Socket) IsClosed() bool {
	socket.Lock()
	defer socket.Unlock()
	return socket.closed
}

// is there data to write?
func (socket *Socket) readyToWrite() bool {
	socket.Lock()
	defer socket.Unlock()
	// on the first time observing socket.closed, we still have to write socket.finalData
	return !socket.finalized && (socket.totalLength > 0 || socket.closed)
}

// send actually writes messages to socket.Conn; it may block
func (socket *Socket) send() {
	for {
		// we are holding the trylock: actually do the write
		socket.performWrite()
		// surrender the trylock, avoiding a race where a write comes in after we've
		// checked readyToWrite() and it returned false, but while we still hold the trylock:
		socket.writerSemaphore.Release()
		// check if more data came in while we held the trylock:
		if !socket.readyToWrite() {
			return
		}
		if !socket.writerSemaphore.TryAcquire() {
			// failed to acquire; exit and wait for the holder to observe readyToWrite()
			// after releasing it
			return
		}
		// got the lock again, loop back around and write
	}
}

// write the contents of the buffer, then see if we need to close
// returns whether we closed
func (socket *Socket) performWrite() (closed bool) {
	// retrieve the buffered data, clear the buffer
	socket.Lock()
	buffers := socket.buffers
	socket.buffers = nil
	socket.totalLength = 0
	closed = socket.closed
	socket.Unlock()

	var err error
	if 0 < len(buffers) {
		socket.conn.WriteBuffers(buffers)
	}

	closed = closed || err != nil
	if closed {
		socket.finalize()
	}
	return
}

// mark closed and send final data. you must be holding the semaphore to call this:
func (socket *Socket) finalize() {
	// mark the socket closed (if someone hasn't already), then write error lines
	socket.Lock()
	socket.closed = true
	finalized := socket.finalized
	socket.finalized = true
	finalData := socket.finalData
	if socket.sendQExceeded {
		finalData = sendQExceededMessage
	}
	socket.Unlock()

	if finalized {
		return
	}

	if len(finalData) != 0 {
		socket.conn.Write(finalData)
	}

	// close the connection
	socket.conn.Close()
}
<commit_msg>fix missing err check, thanks @eklitzke<commit_after>// Copyright (c) 2012-2014 Jeremy Latt
// Copyright (c) 2016-2017 Daniel Oaks <daniel@danieloaks.net>
// released under the MIT license

package irc

import (
	"errors"
	"io"
	"strings"
	"sync"

	"github.com/oragono/oragono/irc/utils"
)

var (
	errSendQExceeded = errors.New("SendQ exceeded")

	sendQExceededMessage = []byte("\r\nERROR :SendQ Exceeded\r\n")
)

// Socket represents an IRC socket.
type Socket struct {
	sync.Mutex

	conn IRCConn

	maxSendQBytes int

	// this is a trylock enforcing that only one goroutine can write to `conn` at a time
	writerSemaphore utils.Semaphore

	buffers       [][]byte
	totalLength   int
	closed        bool
	sendQExceeded bool
	finalData     []byte // what to send when we die
	finalized     bool
}

// NewSocket returns a new Socket.
func NewSocket(conn IRCConn, maxSendQBytes int) *Socket {
	result := Socket{
		conn:          conn,
		maxSendQBytes: maxSendQBytes,
	}
	result.writerSemaphore.Initialize(1)
	return &result
}

// Close stops a Socket from being able to send/receive any more data.
func (socket *Socket) Close() {
	socket.Lock()
	socket.closed = true
	socket.Unlock()

	socket.wakeWriter()
}

// Read returns a single IRC line from a Socket.
func (socket *Socket) Read() (string, error) {
	if socket.IsClosed() {
		return "", io.EOF
	}

	lineBytes, err := socket.conn.ReadLine()

	// convert bytes to string
	line := string(lineBytes)

	// read last message properly (such as ERROR/QUIT/etc), just fail next reads/writes
	if err == io.EOF {
		socket.Close()
	}

	if err == io.EOF && strings.TrimSpace(line) != "" {
		// don't do anything
	} else if err != nil {
		return "", err
	}

	return line, nil
}

// Write sends the given string out of Socket. Requirements:
// 1. MUST NOT block for macroscopic amounts of time
// 2. MUST NOT reorder messages
// 3. MUST provide mutual exclusion for socket.conn.Write
// 4. SHOULD NOT tie up additional goroutines, beyond the one blocked on socket.conn.Write
func (socket *Socket) Write(data []byte) (err error) {
	if len(data) == 0 {
		return
	}

	socket.Lock()
	if socket.closed {
		err = io.EOF
	} else {
		prospectiveLen := socket.totalLength + len(data)
		if prospectiveLen > socket.maxSendQBytes {
			socket.sendQExceeded = true
			socket.closed = true
			err = errSendQExceeded
		} else {
			socket.buffers = append(socket.buffers, data)
			socket.totalLength = prospectiveLen
		}
	}
	socket.Unlock()

	socket.wakeWriter()
	return
}

// BlockingWrite sends the given string out of Socket. Requirements:
// 1. MUST block until the message is sent
// 2. MUST bypass sendq (calls to BlockingWrite cannot, on their own, cause a sendq overflow)
// 3. MUST provide mutual exclusion for socket.conn.Write
// 4. MUST respect the same ordering guarantees as Write (i.e., if a call to Write that sends
//    message m1 happens-before a call to BlockingWrite that sends message m2,
//    m1 must be sent on the wire before m2
// Callers MUST be writing to the client's socket from the client's own goroutine;
// other callers must use the nonblocking Write call instead. Otherwise, a client
// with a slow/unreliable connection risks stalling the progress of the system as a whole.
func (socket *Socket) BlockingWrite(data []byte) (err error) {
	if len(data) == 0 {
		return
	}

	// after releasing the semaphore, we must check for fresh data, same as `send`
	defer func() {
		if socket.readyToWrite() {
			socket.wakeWriter()
		}
	}()

	// blocking acquire of the trylock
	socket.writerSemaphore.Acquire()
	defer socket.writerSemaphore.Release()

	// first, flush any buffered data, to preserve the ordering guarantees
	closed := socket.performWrite()
	if closed {
		return io.EOF
	}

	err = socket.conn.Write(data)
	if err != nil {
		socket.finalize()
	}
	return
}

// wakeWriter starts the goroutine that actually performs the write, without blocking
func (socket *Socket) wakeWriter() {
	if socket.writerSemaphore.TryAcquire() {
		// acquired the trylock; send() will release it
		go socket.send()
	}
	// else: do nothing, the holder will check for more data after releasing it
}

// SetFinalData sets the final data to send when the SocketWriter closes.
func (socket *Socket) SetFinalData(data []byte) {
	socket.Lock()
	defer socket.Unlock()
	socket.finalData = data
}

// IsClosed returns whether the socket is closed.
func (socket *Socket) IsClosed() bool {
	socket.Lock()
	defer socket.Unlock()
	return socket.closed
}

// is there data to write?
func (socket *Socket) readyToWrite() bool {
	socket.Lock()
	defer socket.Unlock()
	// on the first time observing socket.closed, we still have to write socket.finalData
	return !socket.finalized && (socket.totalLength > 0 || socket.closed)
}

// send actually writes messages to socket.Conn; it may block
func (socket *Socket) send() {
	for {
		// we are holding the trylock: actually do the write
		socket.performWrite()
		// surrender the trylock, avoiding a race where a write comes in after we've
		// checked readyToWrite() and it returned false, but while we still hold the trylock:
		socket.writerSemaphore.Release()
		// check if more data came in while we held the trylock:
		if !socket.readyToWrite() {
			return
		}
		if !socket.writerSemaphore.TryAcquire() {
			// failed to acquire; exit and wait for the holder to observe readyToWrite()
			// after releasing it
			return
		}
		// got the lock again, loop back around and write
	}
}

// write the contents of the buffer, then see if we need to close
// returns whether we closed
func (socket *Socket) performWrite() (closed bool) {
	// retrieve the buffered data, clear the buffer
	socket.Lock()
	buffers := socket.buffers
	socket.buffers = nil
	socket.totalLength = 0
	closed = socket.closed
	socket.Unlock()

	var err error
	if 0 < len(buffers) {
		err = socket.conn.WriteBuffers(buffers)
	}

	closed = closed || err != nil
	if closed {
		socket.finalize()
	}
	return
}

// mark closed and send final data. you must be holding the semaphore to call this:
func (socket *Socket) finalize() {
	// mark the socket closed (if someone hasn't already), then write error lines
	socket.Lock()
	socket.closed = true
	finalized := socket.finalized
	socket.finalized = true
	finalData := socket.finalData
	if socket.sendQExceeded {
		finalData = sendQExceededMessage
	}
	socket.Unlock()

	if finalized {
		return
	}

	if len(finalData) != 0 {
		socket.conn.Write(finalData)
	}

	// close the connection
	socket.conn.Close()
}
<|endoftext|>


<commit_before>package redis

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/bsm/ratelimit.v1"
)

var (
	errClosed      = errors.New("redis: client is closed")
	errPoolTimeout = errors.New("redis: connection pool timeout")
)

type pool interface {
	First() *conn
	Get() (*conn, error)
	Put(*conn) error
	Remove(*conn) error
	Len() int
	FreeLen() int
	Close() error
}

type connList struct {
	cns  []*conn
	mx   sync.Mutex
	len  int32 // atomic
	size int32
}

func newConnList(size int) *connList {
	return &connList{
		cns:  make([]*conn, 0, size),
		size: int32(size),
	}
}

func (l *connList) Len() int {
	return int(atomic.LoadInt32(&l.len))
}

// Reserve reserves place in the list and returns true on success. The
// caller must add or remove connection if place was reserved.
func (l *connList) Reserve() bool {
	len := atomic.AddInt32(&l.len, 1)
	reserved := len <= l.size
	if !reserved {
		atomic.AddInt32(&l.len, -1)
	}
	return reserved
}

// Add adds connection to the list. The caller must reserve place first.
func (l *connList) Add(cn *conn) {
	l.mx.Lock()
	l.cns = append(l.cns, cn)
	l.mx.Unlock()
}

// Remove closes connection and removes it from the list.
func (l *connList) Remove(cn *conn) error {
	defer l.mx.Unlock()
	l.mx.Lock()

	if cn == nil {
		atomic.AddInt32(&l.len, -1)
		return nil
	}

	for i, c := range l.cns {
		if c == cn {
			l.cns = append(l.cns[:i], l.cns[i+1:]...)
			atomic.AddInt32(&l.len, -1)
			return cn.Close()
		}
	}

	if l.closed() {
		return nil
	}
	panic("conn not found in the list")
}

func (l *connList) Replace(cn, newcn *conn) error {
	defer l.mx.Unlock()
	l.mx.Lock()

	for i, c := range l.cns {
		if c == cn {
			l.cns[i] = newcn
			return cn.Close()
		}
	}

	if l.closed() {
		return newcn.Close()
	}
	panic("conn not found in the list")
}

func (l *connList) Close() (retErr error) {
	l.mx.Lock()
	for _, c := range l.cns {
		if err := c.Close(); err != nil {
			retErr = err
		}
	}
	l.cns = nil
	atomic.StoreInt32(&l.len, 0)
	l.mx.Unlock()
	return retErr
}

func (l *connList) closed() bool {
	return l.cns == nil
}

type connPool struct {
	dialer func() (*conn, error)

	rl        *ratelimit.RateLimiter
	opt       *Options
	conns     *connList
	freeConns chan *conn

	_closed int32

	lastDialErr error
}

func newConnPool(opt *Options) *connPool {
	p := &connPool{
		dialer: newConnDialer(opt),

		rl:        ratelimit.New(2*opt.getPoolSize(), time.Second),
		opt:       opt,
		conns:     newConnList(opt.getPoolSize()),
		freeConns: make(chan *conn, opt.getPoolSize()),
	}
	if p.opt.getIdleTimeout() > 0 {
		go p.reaper()
	}
	return p
}

func (p *connPool) closed() bool {
	return atomic.LoadInt32(&p._closed) == 1
}

func (p *connPool) isIdle(cn *conn) bool {
	return p.opt.getIdleTimeout() > 0 && time.Since(cn.usedAt) > p.opt.getIdleTimeout()
}

// First returns first non-idle connection from the pool or nil if
// there are no connections.
func (p *connPool) First() *conn {
	for {
		select {
		case cn := <-p.freeConns:
			if p.isIdle(cn) {
				p.conns.Remove(cn)
				continue
			}
			return cn
		default:
			return nil
		}
	}
	panic("not reached")
}

// wait waits for free non-idle connection. It returns nil on timeout.
func (p *connPool) wait() *conn {
	deadline := time.After(p.opt.getPoolTimeout())
	for {
		select {
		case cn := <-p.freeConns:
			if p.isIdle(cn) {
				p.Remove(cn)
				continue
			}
			return cn
		case <-deadline:
			return nil
		}
	}
	panic("not reached")
}

// Establish a new connection
func (p *connPool) new() (*conn, error) {
	if p.rl.Limit() {
		err := fmt.Errorf(
			"redis: you open connections too fast (last error: %v)",
			p.lastDialErr,
		)
		return nil, err
	}

	cn, err := p.dialer()
	if err != nil {
		p.lastDialErr = err
		return nil, err
	}

	return cn, nil
}

// Get returns existed connection from the pool or creates a new one.
func (p *connPool) Get() (*conn, error) {
	if p.closed() {
		return nil, errClosed
	}

	// Fetch first non-idle connection, if available.
	if cn := p.First(); cn != nil {
		return cn, nil
	}

	// Try to create a new one.
	if p.conns.Reserve() {
		cn, err := p.new()
		if err != nil {
			p.conns.Remove(nil)
			return nil, err
		}
		p.conns.Add(cn)
		return cn, nil
	}

	// Otherwise, wait for the available connection.
	if cn := p.wait(); cn != nil {
		return cn, nil
	}

	return nil, errPoolTimeout
}

func (p *connPool) Put(cn *conn) error {
	if cn.rd.Buffered() != 0 {
		b, _ := cn.rd.ReadN(cn.rd.Buffered())
		log.Printf("redis: connection has unread data: %q", b)
		return p.Remove(cn)
	}
	if p.opt.getIdleTimeout() > 0 {
		cn.usedAt = time.Now()
	}
	p.freeConns <- cn
	return nil
}

func (p *connPool) Remove(cn *conn) error {
	// Replace existing connection with new one and unblock waiter.
	newcn, err := p.new()
	if err != nil {
		return p.conns.Remove(cn)
	}
	p.freeConns <- newcn
	return p.conns.Replace(cn, newcn)
}

// Len returns total number of connections.
func (p *connPool) Len() int {
	return p.conns.Len()
}

// FreeLen returns number of free connections.
func (p *connPool) FreeLen() int {
	return len(p.freeConns)
}

func (p *connPool) Close() (retErr error) {
	if !atomic.CompareAndSwapInt32(&p._closed, 0, 1) {
		return errClosed
	}
	// First close free connections.
	for p.Len() > 0 {
		cn := p.wait()
		if cn == nil {
			break
		}
		if err := p.conns.Remove(cn); err != nil {
			retErr = err
		}
	}
	// Then close the rest.
	if err := p.conns.Close(); err != nil {
		retErr = err
	}
	return retErr
}

func (p *connPool) reaper() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for _ = range ticker.C {
		if p.closed() {
			break
		}

		// pool.First removes idle connections from the pool and
		// returns first non-idle connection. So just put returned
		// connection back.
		if cn := p.First(); cn != nil {
			p.Put(cn)
		}
	}
}

//------------------------------------------------------------------------------

type singleConnPool struct {
	pool pool

	cnMtx sync.Mutex
	cn    *conn

	reusable bool

	closed bool
}

func newSingleConnPool(pool pool, reusable bool) *singleConnPool {
	return &singleConnPool{
		pool:     pool,
		reusable: reusable,
	}
}

func (p *singleConnPool) SetConn(cn *conn) {
	p.cnMtx.Lock()
	p.cn = cn
	p.cnMtx.Unlock()
}

func (p *singleConnPool) First() *conn {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	return p.cn
}

func (p *singleConnPool) Get() (*conn, error) {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()

	if p.closed {
		return nil, errClosed
	}
	if p.cn != nil {
		return p.cn, nil
	}

	cn, err := p.pool.Get()
	if err != nil {
		return nil, err
	}
	p.cn = cn

	return p.cn, nil
}

func (p *singleConnPool) Put(cn *conn) error {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn != cn {
		panic("p.cn != cn")
	}
	if p.closed {
		return errClosed
	}
	return nil
}

func (p *singleConnPool) Remove(cn *conn) error {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn == nil {
		panic("p.cn == nil")
	}
	if p.cn != cn {
		panic("p.cn != cn")
	}
	if p.closed {
		return errClosed
	}
	return p.remove()
}

func (p *singleConnPool) remove() error {
	err := p.pool.Remove(p.cn)
	p.cn = nil
	return err
}

func (p *singleConnPool) Len() int {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn == nil {
		return 0
	}
	return 1
}

func (p *singleConnPool) FreeLen() int {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn == nil {
		return 0
	}
	return 1
}

func (p *singleConnPool) Close() error {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.closed {
		return nil
	}
	p.closed = true
	var err error
	if p.cn != nil {
		if p.reusable {
			err = p.pool.Put(p.cn)
			p.cn = nil
		} else {
			err = p.remove()
		}
	}
	return err
}
<commit_msg>pool: close all connections at once.<commit_after>package redis

import (
	"errors"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/bsm/ratelimit.v1"
)

var (
	errClosed      = errors.New("redis: client is closed")
	errPoolTimeout = errors.New("redis: connection pool timeout")
)

type pool interface {
	First() *conn
	Get() (*conn, error)
	Put(*conn) error
	Remove(*conn) error
	Len() int
	FreeLen() int
	Close() error
}

type connList struct {
	cns  []*conn
	mx   sync.Mutex
	len  int32 // atomic
	size int32
}

func newConnList(size int) *connList {
	return &connList{
		cns:  make([]*conn, 0, size),
		size: int32(size),
	}
}

func (l *connList) Len() int {
	return int(atomic.LoadInt32(&l.len))
}

// Reserve reserves place in the list and returns true on success. The
// caller must add or remove connection if place was reserved.
func (l *connList) Reserve() bool {
	len := atomic.AddInt32(&l.len, 1)
	reserved := len <= l.size
	if !reserved {
		atomic.AddInt32(&l.len, -1)
	}
	return reserved
}

// Add adds connection to the list. The caller must reserve place first.
func (l *connList) Add(cn *conn) {
	l.mx.Lock()
	l.cns = append(l.cns, cn)
	l.mx.Unlock()
}

// Remove closes connection and removes it from the list.
func (l *connList) Remove(cn *conn) error {
	defer l.mx.Unlock()
	l.mx.Lock()

	if cn == nil {
		atomic.AddInt32(&l.len, -1)
		return nil
	}

	for i, c := range l.cns {
		if c == cn {
			l.cns = append(l.cns[:i], l.cns[i+1:]...)
			atomic.AddInt32(&l.len, -1)
			return cn.Close()
		}
	}

	if l.closed() {
		return nil
	}
	panic("conn not found in the list")
}

func (l *connList) Replace(cn, newcn *conn) error {
	defer l.mx.Unlock()
	l.mx.Lock()

	for i, c := range l.cns {
		if c == cn {
			l.cns[i] = newcn
			return cn.Close()
		}
	}

	if l.closed() {
		return newcn.Close()
	}
	panic("conn not found in the list")
}

func (l *connList) Close() (retErr error) {
	l.mx.Lock()
	for _, c := range l.cns {
		if err := c.Close(); err != nil {
			retErr = err
		}
	}
	l.cns = nil
	atomic.StoreInt32(&l.len, 0)
	l.mx.Unlock()
	return retErr
}

func (l *connList) closed() bool {
	return l.cns == nil
}

type connPool struct {
	dialer func() (*conn, error)

	rl        *ratelimit.RateLimiter
	opt       *Options
	conns     *connList
	freeConns chan *conn

	_closed int32

	lastDialErr error
}

func newConnPool(opt *Options) *connPool {
	p := &connPool{
		dialer: newConnDialer(opt),

		rl:        ratelimit.New(2*opt.getPoolSize(), time.Second),
		opt:       opt,
		conns:     newConnList(opt.getPoolSize()),
		freeConns: make(chan *conn, opt.getPoolSize()),
	}
	if p.opt.getIdleTimeout() > 0 {
		go p.reaper()
	}
	return p
}

func (p *connPool) closed() bool {
	return atomic.LoadInt32(&p._closed) == 1
}

func (p *connPool) isIdle(cn *conn) bool {
	return p.opt.getIdleTimeout() > 0 && time.Since(cn.usedAt) > p.opt.getIdleTimeout()
}

// First returns first non-idle connection from the pool or nil if
// there are no connections.
func (p *connPool) First() *conn {
	for {
		select {
		case cn := <-p.freeConns:
			if p.isIdle(cn) {
				p.conns.Remove(cn)
				continue
			}
			return cn
		default:
			return nil
		}
	}
	panic("not reached")
}

// wait waits for free non-idle connection. It returns nil on timeout.
func (p *connPool) wait() *conn {
	deadline := time.After(p.opt.getPoolTimeout())
	for {
		select {
		case cn := <-p.freeConns:
			if p.isIdle(cn) {
				p.Remove(cn)
				continue
			}
			return cn
		case <-deadline:
			return nil
		}
	}
	panic("not reached")
}

// Establish a new connection
func (p *connPool) new() (*conn, error) {
	if p.rl.Limit() {
		err := fmt.Errorf(
			"redis: you open connections too fast (last error: %v)",
			p.lastDialErr,
		)
		return nil, err
	}

	cn, err := p.dialer()
	if err != nil {
		p.lastDialErr = err
		return nil, err
	}

	return cn, nil
}

// Get returns existed connection from the pool or creates a new one.
func (p *connPool) Get() (*conn, error) {
	if p.closed() {
		return nil, errClosed
	}

	// Fetch first non-idle connection, if available.
	if cn := p.First(); cn != nil {
		return cn, nil
	}

	// Try to create a new one.
	if p.conns.Reserve() {
		cn, err := p.new()
		if err != nil {
			p.conns.Remove(nil)
			return nil, err
		}
		p.conns.Add(cn)
		return cn, nil
	}

	// Otherwise, wait for the available connection.
	if cn := p.wait(); cn != nil {
		return cn, nil
	}

	return nil, errPoolTimeout
}

func (p *connPool) Put(cn *conn) error {
	if cn.rd.Buffered() != 0 {
		b, _ := cn.rd.ReadN(cn.rd.Buffered())
		log.Printf("redis: connection has unread data: %q", b)
		return p.Remove(cn)
	}
	if p.opt.getIdleTimeout() > 0 {
		cn.usedAt = time.Now()
	}
	p.freeConns <- cn
	return nil
}

func (p *connPool) Remove(cn *conn) error {
	// Replace existing connection with new one and unblock waiter.
	newcn, err := p.new()
	if err != nil {
		return p.conns.Remove(cn)
	}
	p.freeConns <- newcn
	return p.conns.Replace(cn, newcn)
}

// Len returns total number of connections.
func (p *connPool) Len() int {
	return p.conns.Len()
}

// FreeLen returns number of free connections.
func (p *connPool) FreeLen() int {
	return len(p.freeConns)
}

func (p *connPool) Close() (retErr error) {
	if !atomic.CompareAndSwapInt32(&p._closed, 0, 1) {
		return errClosed
	}
	// Wait for app to free connections, but don't close them immediately.
	for i := 0; i < p.Len(); i++ {
		if cn := p.wait(); cn == nil {
			break
		}
	}
	// Close all connections.
	if err := p.conns.Close(); err != nil {
		retErr = err
	}
	return retErr
}

func (p *connPool) reaper() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for _ = range ticker.C {
		if p.closed() {
			break
		}

		// pool.First removes idle connections from the pool and
		// returns first non-idle connection. So just put returned
		// connection back.
		if cn := p.First(); cn != nil {
			p.Put(cn)
		}
	}
}

//------------------------------------------------------------------------------

type singleConnPool struct {
	pool pool

	cnMtx sync.Mutex
	cn    *conn

	reusable bool

	closed bool
}

func newSingleConnPool(pool pool, reusable bool) *singleConnPool {
	return &singleConnPool{
		pool:     pool,
		reusable: reusable,
	}
}

func (p *singleConnPool) SetConn(cn *conn) {
	p.cnMtx.Lock()
	p.cn = cn
	p.cnMtx.Unlock()
}

func (p *singleConnPool) First() *conn {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	return p.cn
}

func (p *singleConnPool) Get() (*conn, error) {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()

	if p.closed {
		return nil, errClosed
	}
	if p.cn != nil {
		return p.cn, nil
	}

	cn, err := p.pool.Get()
	if err != nil {
		return nil, err
	}
	p.cn = cn

	return p.cn, nil
}

func (p *singleConnPool) Put(cn *conn) error {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn != cn {
		panic("p.cn != cn")
	}
	if p.closed {
		return errClosed
	}
	return nil
}

func (p *singleConnPool) Remove(cn *conn) error {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn == nil {
		panic("p.cn == nil")
	}
	if p.cn != cn {
		panic("p.cn != cn")
	}
	if p.closed {
		return errClosed
	}
	return p.remove()
}

func (p *singleConnPool) remove() error {
	err := p.pool.Remove(p.cn)
	p.cn = nil
	return err
}

func (p *singleConnPool) Len() int {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn == nil {
		return 0
	}
	return 1
}

func (p *singleConnPool) FreeLen() int {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.cn == nil {
		return 0
	}
	return 1
}

func (p *singleConnPool) Close() error {
	defer p.cnMtx.Unlock()
	p.cnMtx.Lock()
	if p.closed {
		return nil
	}
	p.closed = true
	var err error
	if p.cn != nil {
		if p.reusable {
			err = p.pool.Put(p.cn)
			p.cn = nil
		} else {
			err = p.remove()
		}
	}
	return err
}
<|endoftext|>


<commit_before>package edc

import (
	"context"
	"time"
)

// Post - struct for post
type Post struct {
	ID        int64  `sql:"id"         json:"id"   form:"id"   query:"id"`
	Name      string `sql:"name"       json:"name" form:"name" query:"name"`
	GO        bool   `sql:"go"         json:"go"   form:"go"   query:"go"`
	Note      string `sql:"note"       json:"note" form:"note" query:"note"`
	CreatedAt string `sql:"created_at" json:"-"`
	UpdatedAt string `sql:"updated_at" json:"-"`
}

// PostList - struct for post list
type PostList struct {
	ID   int64  `sql:"id"   json:"id"   form:"id"   query:"id"`
	Name string `sql:"name" json:"name" form:"name" query:"name"`
	GO   bool   `sql:"go"   json:"go"   form:"go"   query:"go"`
	Note string `sql:"note" json:"note" form:"note" query:"note"`
}

// PostGet - get one post by id
func PostGet(id int64) (Post, error) {
	var post Post
	if id == 0 {
		return post, nil
	}
	post.ID = id
	err := pool.QueryRow(context.Background(), `
		SELECT
			name,
			go,
			note,
			created_at,
			updated_at
		FROM
			kinds
		WHERE
			id = $1
	`, id).Scan(&post.Name, &post.GO, &post.Note, &post.CreatedAt, &post.UpdatedAt)
	if err != nil {
		errmsg("PostGet QueryRow", err)
	}
	return post, nil
}

// PostListGet - get all post for list
func PostListGet() ([]PostList, error) {
	var posts []PostList
	rows, err := pool.Query(context.Background(), `
		SELECT
			id,
			name,
			go,
			note,
		FROM
			kinds
	`)
	if err != nil {
		errmsg("PostListGet Query", err)
	}
	for rows.Next() {
		var post PostList
		err := rows.Scan(&post.ID, &post.Name, &post.GO, &post.Note)
		if err != nil {
			errmsg("PostListGet Scan", err)
			return posts, err
		}
		posts = append(posts, post)
	}
	return posts, rows.Err()
}

// PostSelectGet - get all post for select
func PostSelectGet(g bool) ([]SelectItem, error) {
	var posts []SelectItem
	rows, err := pool.Query(context.Background(), `
		SELECT
			id,
			name
		FROM
			posts
		WHERE
			go = $1
		ORDER BY
			name ASC
	`, g)
	if err != nil {
		errmsg("PostSelectGet Query", err)
	}
	for rows.Next() {
		var post SelectItem
		err := rows.Scan(&post.ID, &post.Name)
		if err != nil {
			errmsg("PostSelectGet Scan", err)
			return posts, err
		}
		posts = append(posts, post)
	}
	return posts, rows.Err()
}

// PostInsert - create new post
func PostInsert(post Post) (int64, error) {
	err := pool.QueryRow(context.Background(), `
		INSERT INTO posts
		(
			name,
			go,
			note,
			created_at,
			updated_at
		)
		VALUES
		(
			$1,
			$2,
			$3,
			$4,
			$5
		)
		RETURNING
			id
	`, post.Name, post.GO, post.Note, time.Now(), time.Now()).Scan(&post.ID)
	if err != nil {
		errmsg("PostInsert QueryRow", err)
	}
	return post.ID, nil
}

// PostUpdate - save post changes
func PostUpdate(post Post) error {
	_, err := pool.Exec(context.Background(), `
		UPDATE posts SET
			name = $2,
			go = $3,
			note = $4,
			updated_at = $5
		WHERE
			id = $1
	`, post.ID, post.Name, post.GO, post.Note, time.Now())
	if err != nil {
		errmsg("UpdatePost update", err)
	}
	return err
}

// PostDelete - delete post by id
func PostDelete(id int64) error {
	if id == 0 {
		return nil
	}
	_, err := pool.Exec(context.Background(), `
		DELETE FROM
			posts
		WHERE
			id = $1
	`, id)
	if err != nil {
		errmsg("DeletePost Exec", err)
	}
	return err
}

func postCreateTable() error {
	str := `
		CREATE TABLE IF NOT EXISTS
			posts (
				id BIGSERIAL PRIMARY KEY,
				name TEXT,
				go BOOL NOT NULL DEFAULT FALSE,
				note TEXT,
				created_at TIMESTAMP without time zone,
				updated_at TIMESTAMP without time zone default now(),
				UNIQUE (name, go)
			)
	`
	_, err := pool.Exec(context.Background(), str)
	if err != nil {
		errmsg("postCreateTable exec", err)
	}
	return err
}
<commit_msg>add missing order<commit_after>package edc

import (
	"context"
	"time"
)

// Post - struct for post
type Post struct {
	ID        int64  `sql:"id"         json:"id"   form:"id"   query:"id"`
	Name      string `sql:"name"       json:"name" form:"name" query:"name"`
	GO        bool   `sql:"go"         json:"go"   form:"go"   query:"go"`
	Note      string `sql:"note"       json:"note" form:"note" query:"note"`
	CreatedAt string `sql:"created_at" json:"-"`
	UpdatedAt string `sql:"updated_at" json:"-"`
}

// PostList - struct for post list
type PostList struct {
	ID   int64  `sql:"id"   json:"id"   form:"id"   query:"id"`
	Name string `sql:"name" json:"name" form:"name" query:"name"`
	GO   bool   `sql:"go"   json:"go"   form:"go"   query:"go"`
	Note string `sql:"note" json:"note" form:"note" query:"note"`
}

// PostGet - get one post by id
func PostGet(id int64) (Post, error) {
	var post Post
	if id == 0 {
		return post, nil
	}
	post.ID = id
	err := pool.QueryRow(context.Background(), `
		SELECT
			name,
			go,
			note,
			created_at,
			updated_at
		FROM
			kinds
		WHERE
			id = $1
	`, id).Scan(&post.Name, &post.GO, &post.Note, &post.CreatedAt, &post.UpdatedAt)
	if err != nil {
		errmsg("PostGet QueryRow", err)
	}
	return post, nil
}

// PostListGet - get all post for list
func PostListGet() ([]PostList, error) {
	var posts []PostList
	rows, err := pool.Query(context.Background(), `
		SELECT
			id,
			name,
			go,
			note
		FROM
			kinds
		ORDER BY
			name ASC
	`)
	if err != nil {
		errmsg("PostListGet Query", err)
	}
	for rows.Next() {
		var post PostList
		err := rows.Scan(&post.ID, &post.Name, &post.GO, &post.Note)
		if err != nil {
			errmsg("PostListGet Scan", err)
			return posts, err
		}
		posts = append(posts, post)
	}
	return posts, rows.Err()
}

// PostSelectGet - get all post for select
func PostSelectGet(g bool) ([]SelectItem, error) {
	var posts []SelectItem
	rows, err := pool.Query(context.Background(), `
		SELECT
			id,
			name
		FROM
			posts
		WHERE
			go = $1
		ORDER BY
			name ASC
	`, g)
	if err != nil {
		errmsg("PostSelectGet Query", err)
	}
	for rows.Next() {
		var post SelectItem
		err := rows.Scan(&post.ID, &post.Name)
		if err != nil {
			errmsg("PostSelectGet Scan", err)
			return posts, err
		}
		posts = append(posts, post)
	}
	return posts, rows.Err()
}

// PostInsert - create new post
func PostInsert(post Post) (int64, error) {
	err := pool.QueryRow(context.Background(), `
		INSERT INTO posts
		(
			name,
			go,
			note,
			created_at,
			updated_at
		)
		VALUES
		(
			$1,
			$2,
			$3,
			$4,
			$5
		)
		RETURNING
			id
	`, post.Name, post.GO, post.Note, time.Now(), time.Now()).Scan(&post.ID)
	if err != nil {
		errmsg("PostInsert QueryRow", err)
	}
	return post.ID, nil
}

// PostUpdate - save post changes
func PostUpdate(post Post) error {
	_, err := pool.Exec(context.Background(), `
		UPDATE posts SET
			name = $2,
			go = $3,
			note = $4,
			updated_at = $5
		WHERE
			id = $1
	`, post.ID, post.Name, post.GO, post.Note, time.Now())
	if err != nil {
		errmsg("UpdatePost update", err)
	}
	return err
}

// PostDelete - delete post by id
func PostDelete(id int64) error {
	if id == 0 {
		return nil
	}
	_, err := pool.Exec(context.Background(), `
		DELETE FROM
			posts
		WHERE
			id = $1
	`, id)
	if err != nil {
		errmsg("DeletePost Exec", err)
	}
	return err
}

func postCreateTable() error {
	str := `
		CREATE TABLE IF NOT EXISTS
			posts (
				id BIGSERIAL PRIMARY KEY,
				name TEXT,
				go BOOL NOT NULL DEFAULT FALSE,
				note TEXT,
				created_at TIMESTAMP without time zone,
				updated_at TIMESTAMP without time zone default now(),
				UNIQUE (name, go)
			)
	`
	_, err := pool.Exec(context.Background(), str)
	if err != nil {
		errmsg("postCreateTable exec", err)
	}
	return err
}
<|endoftext|>