<commit_before>/*
Copyright 2017 Matthew Lord (mattalord@gmail.com)

WARNING: This is experimental and for demonstration purposes only!

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package main

import (
  "os"
  "log"
  "time"
  "flag"
  "sort"
  "fmt"
  "net/http"
  "io/ioutil"
  "encoding/json"
  // uncomment the next import to add profiling to the binary, available via "/debug/pprof" in the RESTful API
  // see: http://blog.ralch.com/tutorial/golang-performance-and-memory-analysis/
  // _ "net/http/pprof"
  "github.com/mattlord/myarbitratord/group_replication/instances"
)

type MembersByOnlineNodes []instances.Instance
var debug = false

var InfoLog = log.New( os.Stderr,
              "INFO: ",
              log.Ldate|log.Ltime|log.Lshortfile )

var DebugLog = log.New( os.Stderr,
               "DEBUG: ",
               log.Ldate|log.Ltime|log.Lshortfile )

// This is where I'll store all operating status metrics, presented as JSON via the "/stats" HTTP API call
type stats struct {
  Start_time string			`json:"Started"`
  Uptime string				`json:"Uptime"`
  Loops  uint				`json:"Loops"`
  Partitions uint			`json:"Partitions"`
  Current_seed *instances.Instance	`json:"Current Seed Node"`
  Last_view *[]instances.Instance	`json:"Last Membership View"`
}
var mystats = stats{ Start_time: time.Now().Format(time.RFC1123), Loops: 0, Partitions: 0, }

// This will simply note the available API calls
func defaultHandler( httpW http.ResponseWriter, httpR *http.Request ){
  if( debug ){
    DebugLog.Println( "Handling HTTP request without API call." )
  }

  fmt.Fprintf( httpW, "Welcome to the MySQL Arbitrator's RESTful API handler!\n\nThe available API calls are:\n/stats: Provide runtime and operational stats\n" )
}

// This will serve the stats via a simple RESTful API
func statsHandler( httpW http.ResponseWriter, httpR *http.Request ){
  if( debug ){
    DebugLog.Printf( "Handling HTTP request for stats. Current stats are: %+v\n", mystats )
  }

  tval, terr := time.Parse( time.RFC1123, mystats.Start_time )
  if( terr != nil ){
    InfoLog.Printf( "Error parsing time value for stats: %+v\n", terr )
  }
  dval := time.Since( tval )
  mystats.Uptime = dval.String()

  statsJSON, err := json.MarshalIndent( &mystats, "", "    " )

  if( err != nil ){
    InfoLog.Printf( "Error handling HTTP request for stats: %+v\n", err )
  }

  fmt.Fprintf( httpW, "%s", statsJSON )
}

func main(){
  var seed_host string
  var seed_port string
  var mysql_user string
  var mysql_pass string
  var mysql_auth_file string
  type json_mysql_auth struct {
    User string      `json:"user"`
    Password string  `json:"password"`
  }

  http.DefaultServeMux.HandleFunc( "/", defaultHandler )
  http.DefaultServeMux.HandleFunc( "/stats", statsHandler )
  var http_port string = "8099"

  flag.StringVar( &seed_host, "seed_host", "", "IP/Hostname of the seed node used to start monitoring the Group Replication cluster (Required Parameter!)" )
  flag.StringVar( &seed_port, "seed_port", "3306", "Port of the seed node used to start monitoring the Group Replication cluster" )
  flag.BoolVar( &debug, "debug", false, "Execute in debug mode with all debug logging enabled" )
  flag.StringVar( &mysql_user, "mysql_user", "root", "The mysql user account to be used when connecting to any node in the cluster" )
  flag.StringVar( &mysql_pass, "mysql_password", "", "The mysql user account password to be used when connecting to any node in the cluster" )
  flag.StringVar( &mysql_auth_file, "mysql_auth_file", "", "The JSON encoded file containining user and password entities for the mysql account to be used when connecting to any node in the cluster" )
  flag.StringVar( &http_port, "http_port", "8099", "The HTTP port used for the RESTful API" )

  flag.Parse()

  // ToDo: I need to handle the password on the command-line more securely
  //       I need to do some data masking for the processlist

  // A host is required, the default port of 3306 will then be attempted
  if( seed_host == "" ){
    fmt.Fprintf( os.Stderr, "No value specified for required flag: -seed_host\n" )
    fmt.Fprintf( os.Stderr, "Usage of %s:\n", os.Args[0] )
    flag.PrintDefaults()
    os.Exit( 1 )
  }

  // let's start a thread to handle the RESTful API calls
  InfoLog.Printf( "Starting HTTP server for RESTful API on port %s\n", http_port )
  go http.ListenAndServe( ":" + http_port, http.DefaultServeMux )

  if( debug ){
    instances.Debug = true
  }

  if( mysql_auth_file != "" && mysql_pass == "" ){
    if( debug ){
      DebugLog.Printf( "Reading MySQL credentials from file: %s\n", mysql_auth_file )
    }

    jsonfile, err := ioutil.ReadFile( mysql_auth_file )

    if( err != nil ){
      log.Fatal( "Could not read mysql credentials from specified file: " + mysql_auth_file )
    }

    var jsonauth json_mysql_auth
    json.Unmarshal( jsonfile, &jsonauth )

    if( debug ){
      DebugLog.Printf( "Unmarshaled mysql auth file contents: %+v\n", jsonauth )
    }

    mysql_user = jsonauth.User
    mysql_pass = jsonauth.Password

    if( mysql_user == "" || mysql_pass == "" ){
      errstr := "Failed to read user and password from " + mysql_auth_file + ". Ensure that the file contents are in the required format: \n{\n  \"user\": \"myser\",\n  \"password\": \"mypass\"\n}"
      log.Fatal( errstr )
    }

    if( debug ){
      DebugLog.Printf( "Read mysql auth info from file. user: %s, password: %s\n", mysql_user, mysql_pass )
    }
  }

  InfoLog.Println( "Welcome to the MySQL Group Replication Arbitrator!" )

  InfoLog.Printf( "Starting operations from seed node: '%s:%s'\n", seed_host, seed_port )
  seed_node := instances.New( seed_host, seed_port, mysql_user, mysql_pass )
  err := MonitorCluster( seed_node )

  if( err != nil ){
    log.Fatal( err )
    os.Exit( 100 )
  } else {
    os.Exit( 0 )
  }
}

func MonitorCluster( seed_node *instances.Instance ) error {
  loop := true
  var err error
  last_view := []instances.Instance{}

  for( loop == true ){
    mystats.Loops = mystats.Loops + 1
    mystats.Current_seed = seed_node
    mystats.Last_view = &last_view

    // let's check the status of the current seed node
    err = seed_node.Connect()
    defer seed_node.Cleanup()

    if( err != nil || seed_node.Member_state != "ONLINE" ){
      // if we couldn't connect to the current seed node or it's no longer part of the group
      // let's try and get a new seed node from the last known membership view
      InfoLog.Println( "Attempting to get a new seed node..." )

      for i := 0; i < len(last_view); i++ {
        if( seed_node != &last_view[i] ){
          err = last_view[i].Connect()
          if( err == nil && last_view[i].Member_state == "ONLINE" ){
            seed_node = &last_view[i]
            InfoLog.Printf( "Updated seed node! New seed node is: '%s:%s'\n", seed_node.Mysql_host, seed_node.Mysql_port )
            break
          }
        }
      }
    }

    members, err := seed_node.GetMembers()

    if( err != nil || seed_node.Online_participants < 1 ){
      // something is up with our current seed node, let's loop again
      continue
    }

    // save this view in case the seed node is no longer valid next time
    last_view = *members

    quorum, err := seed_node.HasQuorum()

    if( err != nil ){
      // something is up with our current seed node, let's loop again
      continue
    }

    if( debug ){
      DebugLog.Printf( "Seed node details: %+v", seed_node )
    }

    if( quorum ){
      // Let's try and shutdown the nodes NOT in the primary partition if we can reach them from the arbitrator

      for _, member := range *members {
        if( member.Member_state == "ERROR" || member.Member_state == "UNREACHABLE" ){
          InfoLog.Printf( "Shutting down non-healthy node: '%s:%s'\n", member.Mysql_host, member.Mysql_port )

          err = member.Connect()

          if( err != nil ){
            InfoLog.Printf( "Could not connect to '%s:%s' in order to shut it down\n", member.Mysql_host, member.Mysql_port )
          } else {
            err = member.Shutdown()
          }

          if( err != nil ){
            InfoLog.Printf( "Could not shutdown instance: '%s:%s'\n", member.Mysql_host, member.Mysql_port )
          }
        }
      }
    } else {
      // handling other network partitions and split brain scenarios will be much trickier... I'll need to try and
      // contact each member in the last seen view and try to determine which partition should become the
      // primary one. We'll then need to contact 1 node in the new primary partition and explicitly set the new
      // membership with 'set global group_replication_force_members="<node_list>"'. Finally we'll need to try
      // and connect to the nodes on the losing side(s) of the partition and attempt to shutdown the mysqlds

      InfoLog.Println( "Network partition detected! Attempting to handle... " )
      mystats.Partitions = mystats.Partitions + 1

      // does anyone have a quorum? Let's double check before forcing the membership
      primary_partition := false

      for i := 0; i < len(last_view); i++ {
        var err error

        err = last_view[i].Connect()

        if( err == nil ){
          quorum, err = last_view[i].HasQuorum()
          // let's make sure that the Online_participants is up to date
          _, err = last_view[i].GetMembers()
        }

        if( err == nil && quorum ){
          seed_node = &last_view[i]
          primary_partition = true
          break
        }
      }

      // If no one in fact has a quorum, then let's see which partition has the most
      // online/participating/communicating members. The participants in that partition
      // will then be the ones that we use to force the new membership and unlock the cluster

      if( primary_partition == false ){
        InfoLog.Println( "No primary partition found! Attempting to choose and force a new one ... " )

        sort.Sort( MembersByOnlineNodes(last_view) )

        // now the last element in the array is the one to use as it's coordinating with the most nodes
        view_len := len(last_view)-1
        seed_node = &last_view[view_len]

        // *BUT*, if there's no clear winner based on sub-partition size, then we should pick the sub-partition (which
        // can be 1 node) that has executed the most GTIDs
        if( last_view[view_len].Online_participants == last_view[view_len-1].Online_participants ){
          bestmemberpos := view_len
          var bestmembertrxcnt uint64 = 0
          var curtrxcnt uint64 = 0
          bestmembertrxcnt, err = last_view[view_len].TransactionsExecutedCount()

          // let's loop backwards through the array as it's sorted by online participants / partition size now
          // skipping the last one as we already have the info for it
          for i := view_len-1; i >= 0; i-- {
            if( last_view[i].Online_participants == last_view[bestmemberpos].Online_participants ){
              curtrxcnt, err = last_view[i].TransactionsExecutedCount()

              if( curtrxcnt > bestmembertrxcnt ){
                bestmembertrxcnt = curtrxcnt
                bestmemberpos = i
              }
            } else {
              // otherwise we've gone backwards far enough and we have the best option
              break
            }
          }

          seed_node = &last_view[bestmemberpos]
        }

        err = seed_node.Connect()

        if( err != nil ){
          // let's just loop again
          continue
        }

        if( debug ){
          DebugLog.Printf( "Member view sorted by number of online nodes: %+v\n", last_view )
        }

        // let's build a string of '<host>:<port>' combinations that we want to use for the new membership view
        members, _ := seed_node.GetMembers()

        force_member_string := ""

        for i, member := range *members {
          err = member.Connect()

          if( err == nil && member.Member_state == "ONLINE" ){
            if( i != 0 ){
              force_member_string = force_member_string + ","
            }

            force_member_string = force_member_string + member.Mysql_host + ":" + member.Mysql_port
          } else {
            member.Member_state = "SHOOT_ME"
          }
        }

        if( force_member_string != "" ){
          InfoLog.Printf( "Forcing group membership to form new primary partition! Using: '%s'\n", force_member_string )

          err := seed_node.ForceMembers( force_member_string )

          if( err != nil ){
            InfoLog.Printf( "Error forcing group membership: %v\n", err )
          } else {
            // We successfully unblocked the group, now let's try and politely STONITH the nodes in the losing partition
            for _, member := range *members {
              if( member.Member_state == "SHOOT_ME" ){
                member.Shutdown()
              }

              if( err != nil ){
                InfoLog.Printf( "Could not shutdown instance: '%s:%s'\n", member.Mysql_host, member.Mysql_port )
              }
            }
          }
        } else {
          InfoLog.Println( "No valid group membership to force!" )
        }
      }
    }

    if( err != nil ){
      loop = false
    } else {
      time.Sleep( time.Millisecond * 2000 )
    }
  }

  return err
}

// The remaining functions are used to sort our membership slice
// We'll never have a super high number of nodes involved, so a simple bubble sort will suffice
func (a MembersByOnlineNodes) Len() int {
  return len(a)
}

func (a MembersByOnlineNodes) Swap( i, j int ) {
  a[i], a[j] = a[j], a[i]
}

func (a MembersByOnlineNodes) Less( i, j int ) bool {
  return a[i].Online_participants < a[j].Online_participants
}
<commit_msg>Improved log messages when dealing with isolated nodes.<commit_after>/*
Copyright 2017 Matthew Lord (mattalord@gmail.com)

WARNING: This is experimental and for demonstration purposes only!

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package main

import (
  "os"
  "log"
  "time"
  "flag"
  "sort"
  "fmt"
  "net/http"
  "io/ioutil"
  "encoding/json"
  // uncomment the next import to add profiling to the binary, available via "/debug/pprof" in the RESTful API
  // see: http://blog.ralch.com/tutorial/golang-performance-and-memory-analysis/
  // _ "net/http/pprof"
  "github.com/mattlord/myarbitratord/group_replication/instances"
)

type MembersByOnlineNodes []instances.Instance
var debug = false

var InfoLog = log.New( os.Stderr,
              "INFO: ",
              log.Ldate|log.Ltime|log.Lshortfile )

var DebugLog = log.New( os.Stderr,
               "DEBUG: ",
               log.Ldate|log.Ltime|log.Lshortfile )

// This is where I'll store all operating status metrics, presented as JSON via the "/stats" HTTP API call
type stats struct {
  Start_time string			`json:"Started"`
  Uptime string				`json:"Uptime"`
  Loops  uint				`json:"Loops"`
  Partitions uint			`json:"Partitions"`
  Current_seed *instances.Instance	`json:"Current Seed Node"`
  Last_view *[]instances.Instance	`json:"Last Membership View"`
}
var mystats = stats{ Start_time: time.Now().Format(time.RFC1123), Loops: 0, Partitions: 0, }

// This will simply note the available API calls
func defaultHandler( httpW http.ResponseWriter, httpR *http.Request ){
  if( debug ){
    DebugLog.Println( "Handling HTTP request without API call." )
  }

  fmt.Fprintf( httpW, "Welcome to the MySQL Arbitrator's RESTful API handler!\n\nThe available API calls are:\n/stats: Provide runtime and operational stats\n" )
}

// This will serve the stats via a simple RESTful API
func statsHandler( httpW http.ResponseWriter, httpR *http.Request ){
  if( debug ){
    DebugLog.Printf( "Handling HTTP request for stats. Current stats are: %+v\n", mystats )
  }

  tval, terr := time.Parse( time.RFC1123, mystats.Start_time )
  if( terr != nil ){
    InfoLog.Printf( "Error parsing time value for stats: %+v\n", terr )
  }
  dval := time.Since( tval )
  mystats.Uptime = dval.String()

  statsJSON, err := json.MarshalIndent( &mystats, "", "    " )

  if( err != nil ){
    InfoLog.Printf( "Error handling HTTP request for stats: %+v\n", err )
  }

  fmt.Fprintf( httpW, "%s", statsJSON )
}

func main(){
  var seed_host string
  var seed_port string
  var mysql_user string
  var mysql_pass string
  var mysql_auth_file string
  type json_mysql_auth struct {
    User string      `json:"user"`
    Password string  `json:"password"`
  }

  http.DefaultServeMux.HandleFunc( "/", defaultHandler )
  http.DefaultServeMux.HandleFunc( "/stats", statsHandler )
  var http_port string = "8099"

  flag.StringVar( &seed_host, "seed_host", "", "IP/Hostname of the seed node used to start monitoring the Group Replication cluster (Required Parameter!)" )
  flag.StringVar( &seed_port, "seed_port", "3306", "Port of the seed node used to start monitoring the Group Replication cluster" )
  flag.BoolVar( &debug, "debug", false, "Execute in debug mode with all debug logging enabled" )
  flag.StringVar( &mysql_user, "mysql_user", "root", "The mysql user account to be used when connecting to any node in the cluster" )
  flag.StringVar( &mysql_pass, "mysql_password", "", "The mysql user account password to be used when connecting to any node in the cluster" )
  flag.StringVar( &mysql_auth_file, "mysql_auth_file", "", "The JSON encoded file containining user and password entities for the mysql account to be used when connecting to any node in the cluster" )
  flag.StringVar( &http_port, "http_port", "8099", "The HTTP port used for the RESTful API" )

  flag.Parse()

  // ToDo: I need to handle the password on the command-line more securely
  //       I need to do some data masking for the processlist

  // A host is required, the default port of 3306 will then be attempted
  if( seed_host == "" ){
    fmt.Fprintf( os.Stderr, "No value specified for required flag: -seed_host\n" )
    fmt.Fprintf( os.Stderr, "Usage of %s:\n", os.Args[0] )
    flag.PrintDefaults()
    os.Exit( 1 )
  }

  // let's start a thread to handle the RESTful API calls
  InfoLog.Printf( "Starting HTTP server for RESTful API on port %s\n", http_port )
  go http.ListenAndServe( ":" + http_port, http.DefaultServeMux )

  if( debug ){
    instances.Debug = true
  }

  if( mysql_auth_file != "" && mysql_pass == "" ){
    if( debug ){
      DebugLog.Printf( "Reading MySQL credentials from file: %s\n", mysql_auth_file )
    }

    jsonfile, err := ioutil.ReadFile( mysql_auth_file )

    if( err != nil ){
      log.Fatal( "Could not read mysql credentials from specified file: " + mysql_auth_file )
    }

    var jsonauth json_mysql_auth
    json.Unmarshal( jsonfile, &jsonauth )

    if( debug ){
      DebugLog.Printf( "Unmarshaled mysql auth file contents: %+v\n", jsonauth )
    }

    mysql_user = jsonauth.User
    mysql_pass = jsonauth.Password

    if( mysql_user == "" || mysql_pass == "" ){
      errstr := "Failed to read user and password from " + mysql_auth_file + ". Ensure that the file contents are in the required format: \n{\n  \"user\": \"myser\",\n  \"password\": \"mypass\"\n}"
      log.Fatal( errstr )
    }

    if( debug ){
      DebugLog.Printf( "Read mysql auth info from file. user: %s, password: %s\n", mysql_user, mysql_pass )
    }
  }

  InfoLog.Println( "Welcome to the MySQL Group Replication Arbitrator!" )

  InfoLog.Printf( "Starting operations from seed node: '%s:%s'\n", seed_host, seed_port )
  seed_node := instances.New( seed_host, seed_port, mysql_user, mysql_pass )
  err := MonitorCluster( seed_node )

  if( err != nil ){
    log.Fatal( err )
    os.Exit( 100 )
  } else {
    os.Exit( 0 )
  }
}

func MonitorCluster( seed_node *instances.Instance ) error {
  loop := true
  var err error
  last_view := []instances.Instance{}

  for( loop == true ){
    mystats.Loops = mystats.Loops + 1
    mystats.Current_seed = seed_node
    mystats.Last_view = &last_view

    // let's check the status of the current seed node
    err = seed_node.Connect()
    defer seed_node.Cleanup()

    if( err != nil || seed_node.Member_state != "ONLINE" ){
      // if we couldn't connect to the current seed node or it's no longer part of the group
      // let's try and get a new seed node from the last known membership view
      InfoLog.Println( "Attempting to get a new seed node..." )

      for i := 0; i < len(last_view); i++ {
        if( seed_node != &last_view[i] ){
          err = last_view[i].Connect()
          if( err == nil && last_view[i].Member_state == "ONLINE" ){
            seed_node = &last_view[i]
            InfoLog.Printf( "Updated seed node! New seed node is: '%s:%s'\n", seed_node.Mysql_host, seed_node.Mysql_port )
            break
          }
        }
      }
    }

    members, err := seed_node.GetMembers()

    if( err != nil || seed_node.Online_participants < 1 ){
      // something is up with our current seed node, let's loop again
      continue
    }

    // save this view in case the seed node is no longer valid next time
    last_view = *members

    quorum, err := seed_node.HasQuorum()

    if( err != nil ){
      // something is up with our current seed node, let's loop again
      continue
    }

    if( debug ){
      DebugLog.Printf( "Seed node details: %+v", seed_node )
    }

    if( quorum ){
      // Let's try and shutdown the nodes NOT in the primary partition if we can reach them from the arbitrator

      for _, member := range *members {
        if( member.Member_state == "ERROR" || member.Member_state == "UNREACHABLE" ){
          InfoLog.Printf( "Shutting down isolated node: '%s:%s'\n", member.Mysql_host, member.Mysql_port )

          err = member.Connect()

          if( err != nil ){
            InfoLog.Printf( "Could not connect to '%s:%s' in order to shut it down\n", member.Mysql_host, member.Mysql_port )
          } else {
            err = member.Shutdown()
          }

          if( err != nil ){
            InfoLog.Printf( "Could not shutdown isolated node: '%s:%s'\n", member.Mysql_host, member.Mysql_port )
          }
        }
      }
    } else {
      // handling other network partitions and split brain scenarios will be much trickier... I'll need to try and
      // contact each member in the last seen view and try to determine which partition should become the
      // primary one. We'll then need to contact 1 node in the new primary partition and explicitly set the new
      // membership with 'set global group_replication_force_members="<node_list>"'. Finally we'll need to try
      // and connect to the nodes on the losing side(s) of the partition and attempt to shutdown the mysqlds

      InfoLog.Println( "Network partition detected! Attempting to handle... " )
      mystats.Partitions = mystats.Partitions + 1

      // does anyone have a quorum? Let's double check before forcing the membership
      primary_partition := false

      for i := 0; i < len(last_view); i++ {
        var err error

        err = last_view[i].Connect()

        if( err == nil ){
          quorum, err = last_view[i].HasQuorum()
          // let's make sure that the Online_participants is up to date
          _, err = last_view[i].GetMembers()
        }

        if( err == nil && quorum ){
          seed_node = &last_view[i]
          primary_partition = true
          break
        }
      }

      // If no one in fact has a quorum, then let's see which partition has the most
      // online/participating/communicating members. The participants in that partition
      // will then be the ones that we use to force the new membership and unlock the cluster

      if( primary_partition == false ){
        InfoLog.Println( "No primary partition found! Attempting to choose and force a new one ... " )

        sort.Sort( MembersByOnlineNodes(last_view) )

        // now the last element in the array is the one to use as it's coordinating with the most nodes
        view_len := len(last_view)-1
        seed_node = &last_view[view_len]

        // *BUT*, if there's no clear winner based on sub-partition size, then we should pick the sub-partition (which
        // can be 1 node) that has executed the most GTIDs
        if( last_view[view_len].Online_participants == last_view[view_len-1].Online_participants ){
          bestmemberpos := view_len
          var bestmembertrxcnt uint64 = 0
          var curtrxcnt uint64 = 0
          bestmembertrxcnt, err = last_view[view_len].TransactionsExecutedCount()

          // let's loop backwards through the array as it's sorted by online participants / partition size now
          // skipping the last one as we already have the info for it
          for i := view_len-1; i >= 0; i-- {
            if( last_view[i].Online_participants == last_view[bestmemberpos].Online_participants ){
              curtrxcnt, err = last_view[i].TransactionsExecutedCount()

              if( curtrxcnt > bestmembertrxcnt ){
                bestmembertrxcnt = curtrxcnt
                bestmemberpos = i
              }
            } else {
              // otherwise we've gone backwards far enough and we have the best option
              break
            }
          }

          seed_node = &last_view[bestmemberpos]
        }

        err = seed_node.Connect()

        if( err != nil ){
          // let's just loop again
          continue
        }

        if( debug ){
          DebugLog.Printf( "Member view sorted by number of online nodes: %+v\n", last_view )
        }

        // let's build a string of '<host>:<port>' combinations that we want to use for the new membership view
        members, _ := seed_node.GetMembers()

        force_member_string := ""

        for i, member := range *members {
          err = member.Connect()

          if( err == nil && member.Member_state == "ONLINE" ){
            if( i != 0 ){
              force_member_string = force_member_string + ","
            }

            force_member_string = force_member_string + member.Mysql_host + ":" + member.Mysql_port
          } else {
            member.Member_state = "SHOOT_ME"
          }
        }

        if( force_member_string != "" ){
          InfoLog.Printf( "Forcing group membership to form new primary partition! Using: '%s'\n", force_member_string )

          err := seed_node.ForceMembers( force_member_string )

          if( err != nil ){
            InfoLog.Printf( "Error forcing group membership: %v\n", err )
          } else {
            // We successfully unblocked the group, now let's try and politely STONITH the nodes in the losing partition
            for _, member := range *members {
              if( member.Member_state == "SHOOT_ME" ){
                member.Shutdown()
              }

              if( err != nil ){
                InfoLog.Printf( "Could not shutdown instance: '%s:%s'\n", member.Mysql_host, member.Mysql_port )
              }
            }
          }
        } else {
          InfoLog.Println( "No valid group membership to force!" )
        }
      }
    }

    if( err != nil ){
      loop = false
    } else {
      time.Sleep( time.Millisecond * 2000 )
    }
  }

  return err
}

// The remaining functions are used to sort our membership slice
// We'll never have a super high number of nodes involved, so a simple bubble sort will suffice
func (a MembersByOnlineNodes) Len() int {
  return len(a)
}

func (a MembersByOnlineNodes) Swap( i, j int ) {
  a[i], a[j] = a[j], a[i]
}

func (a MembersByOnlineNodes) Less( i, j int ) bool {
  return a[i].Online_participants < a[j].Online_participants
}
<|endoftext|>


<commit_before>package email

import (
	"crypto/tls"
	"errors"
	"net"
	"net/mail"
	"net/smtp"
	"net/textproto"
	"sync"
	"syscall"
	"time"
)

type Pool struct {
	addr         string
	auth         smtp.Auth
	max          int
	created      int
	clients      chan *client
	rebuild      chan struct{}
	mut          *sync.Mutex
	lastBuildErr *timestampedErr
	closing      chan struct{}
}

type client struct {
	*smtp.Client
	failCount int
}

type timestampedErr struct {
	err error
	ts  time.Time
}

const maxFails = 4

var (
	ErrClosed  = errors.New("pool closed")
	ErrTimeout = errors.New("timed out")
)

func NewPool(address string, count int, auth smtp.Auth) *Pool {
	return &Pool{
		addr:    address,
		auth:    auth,
		max:     count,
		clients: make(chan *client, count),
		rebuild: make(chan struct{}),
		closing: make(chan struct{}),
		mut:     &sync.Mutex{},
	}
}

// go1.1 didn't have this method
func (c *client) Close() error {
	return c.Text.Close()
}

func (p *Pool) get(timeout time.Duration) *client {
	select {
	case c := <-p.clients:
		return c
	default:
	}

	if p.created < p.max {
		p.makeOne()
	}

	var deadline <-chan time.Time
	if timeout < 0 {
		deadline = nil
	} else {
		deadline = time.After(timeout)
	}

	for {
		select {
		case c := <-p.clients:
			return c
		case <-p.rebuild:
			p.makeOne()
		case <-deadline:
			return nil
		case <-p.closing:
			return nil
		}
	}
}

func shouldReuse(err error) bool {
	// certainly not perfect, but might be close:
	//  - textproto.Errors were valid SMTP over a valid connection,
	//    but resulted from an SMTP error response
	//  - textproto.ProtocolErrors result from connections going down,
	//    invalid SMTP, that sort of thing
	//  - syscall.Errno is probably down connection/bad pipe, but
	//    passed straight through by textproto instead of becoming a
	//    ProtocolError
	//  - if we don't recognize the error, don't reuse the connection
	// A false positive will probably fail on the Reset(), and even if
	// not will eventually hit maxFails.
	// A false negative will knock over (and trigger replacement of) a
	// conn that might have still worked.
	switch err.(type) {
	case *textproto.Error:
		return true
	case *textproto.ProtocolError, textproto.ProtocolError:
		return false
	case syscall.Errno:
		return false
	default:
		return false
	}
}

func (p *Pool) replace(c *client) {
	p.clients <- c
}

func (p *Pool) inc() bool {
	if p.created >= p.max {
		return false
	}

	p.mut.Lock()
	defer p.mut.Unlock()

	if p.created >= p.max {
		return false
	}
	p.created++
	return true
}

func (p *Pool) dec() {
	p.mut.Lock()
	p.created--
	p.mut.Unlock()

	select {
	case p.rebuild <- struct{}{}:
	default:
	}
}

func (p *Pool) makeOne() {
	go func() {
		if p.inc() {
			if c, err := p.build(); err == nil {
				p.clients <- c
			} else {
				p.lastBuildErr = &timestampedErr{err, time.Now()}
				p.dec()
			}
		}
	}()
}

func startTLS(c *smtp.Client, addr string) (bool, error) {
	if ok, _ := c.Extension("STARTTLS"); !ok {
		return false, nil
	}

	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return false, err
	}

	if err := c.StartTLS(&tls.Config{ServerName: host}); err != nil {
		return false, err
	}

	return true, nil
}

func addAuth(c *smtp.Client, auth smtp.Auth) (bool, error) {
	if ok, _ := c.Extension("AUTH"); !ok {
		return false, nil
	}

	if err := c.Auth(auth); err != nil {
		return false, err
	}

	return true, nil
}

func (p *Pool) build() (*client, error) {
	c, err := smtp.Dial(p.addr)
	if err != nil {
		return nil, err
	}

	if _, err := startTLS(c, p.addr); err != nil {
		c.Close()
		return nil, err
	}

	if p.auth != nil {
		if _, err := addAuth(c, p.auth); err != nil {
			c.Close()
			return nil, err
		}
	}

	return &client{c, 0}, nil
}

func (p *Pool) maybeReplace(err error, c *client) {
	if err == nil {
		c.failCount = 0
		p.replace(c)
		return
	}

	c.failCount++
	if c.failCount >= maxFails {
		goto shutdown
	}

	if !shouldReuse(err) {
		goto shutdown
	}

	if err := c.Reset(); err != nil {
		goto shutdown
	}

	p.replace(c)
	return

shutdown:
	p.dec()
	c.Close()
}

func (p *Pool) failedToGet(startTime time.Time) error {
	select {
	case <-p.closing:
		return ErrClosed
	default:
	}

	if p.lastBuildErr != nil && startTime.Before(p.lastBuildErr.ts) {
		return p.lastBuildErr.err
	}

	return ErrTimeout
}

// Send sends an email via a connection pulled from the Pool. The timeout may
// be <0 to indicate no timeout. Otherwise reaching the timeout will produce
// and error building a connection that occurred while we were waiting, or
// otherwise ErrTimeout.
func (p *Pool) Send(e *Email, timeout time.Duration) (err error) {
	start := time.Now()
	c := p.get(timeout)
	if c == nil {
		return p.failedToGet(start)
	}

	defer func() {
		p.maybeReplace(err, c)
	}()

	recipients, err := addressLists(e.To, e.Cc, e.Bcc)
	if err != nil {
		return
	}

	msg, err := e.Bytes()
	if err != nil {
		return
	}

	from, err := emailOnly(e.From)
	if err != nil {
		return
	}
	if err = c.Mail(from); err != nil {
		return
	}

	for _, recip := range recipients {
		if err = c.Rcpt(recip); err != nil {
			return
		}
	}

	w, err := c.Data()
	if err != nil {
		return
	}
	if _, err = w.Write(msg); err != nil {
		return
	}

	err = w.Close()

	return
}

func emailOnly(full string) (string, error) {
	addr, err := mail.ParseAddress(full)
	if err != nil {
		return "", err
	}
	return addr.Address, nil
}

func addressLists(lists ...[]string) ([]string, error) {
	length := 0
	for _, lst := range lists {
		length += len(lst)
	}
	combined := make([]string, 0, length)

	for _, lst := range lists {
		for _, full := range lst {
			addr, err := emailOnly(full)
			if err != nil {
				return nil, err
			}
			combined = append(combined, addr)
		}
	}

	return combined, nil
}

// Close immediately changes the pool's state so no new connections will be
// created, then gets and closes the existing ones as they become available.
func (p *Pool) Close() {
	close(p.closing)

	for p.created > 0 {
		c := <-p.clients
		c.Quit()
		p.dec()
	}
}
<commit_msg>convert to a client wrapper struct earlier, to catch all the Close cases<commit_after>package email

import (
	"crypto/tls"
	"errors"
	"net"
	"net/mail"
	"net/smtp"
	"net/textproto"
	"sync"
	"syscall"
	"time"
)

type Pool struct {
	addr         string
	auth         smtp.Auth
	max          int
	created      int
	clients      chan *client
	rebuild      chan struct{}
	mut          *sync.Mutex
	lastBuildErr *timestampedErr
	closing      chan struct{}
}

type client struct {
	*smtp.Client
	failCount int
}

type timestampedErr struct {
	err error
	ts  time.Time
}

const maxFails = 4

var (
	ErrClosed  = errors.New("pool closed")
	ErrTimeout = errors.New("timed out")
)

func NewPool(address string, count int, auth smtp.Auth) *Pool {
	return &Pool{
		addr:    address,
		auth:    auth,
		max:     count,
		clients: make(chan *client, count),
		rebuild: make(chan struct{}),
		closing: make(chan struct{}),
		mut:     &sync.Mutex{},
	}
}

// go1.1 didn't have this method
func (c *client) Close() error {
	return c.Text.Close()
}

func (p *Pool) get(timeout time.Duration) *client {
	select {
	case c := <-p.clients:
		return c
	default:
	}

	if p.created < p.max {
		p.makeOne()
	}

	var deadline <-chan time.Time
	if timeout < 0 {
		deadline = nil
	} else {
		deadline = time.After(timeout)
	}

	for {
		select {
		case c := <-p.clients:
			return c
		case <-p.rebuild:
			p.makeOne()
		case <-deadline:
			return nil
		case <-p.closing:
			return nil
		}
	}
}

func shouldReuse(err error) bool {
	// certainly not perfect, but might be close:
	//  - textproto.Errors were valid SMTP over a valid connection,
	//    but resulted from an SMTP error response
	//  - textproto.ProtocolErrors result from connections going down,
	//    invalid SMTP, that sort of thing
	//  - syscall.Errno is probably down connection/bad pipe, but
	//    passed straight through by textproto instead of becoming a
	//    ProtocolError
	//  - if we don't recognize the error, don't reuse the connection
	// A false positive will probably fail on the Reset(), and even if
	// not will eventually hit maxFails.
	// A false negative will knock over (and trigger replacement of) a
	// conn that might have still worked.
	switch err.(type) {
	case *textproto.Error:
		return true
	case *textproto.ProtocolError, textproto.ProtocolError:
		return false
	case syscall.Errno:
		return false
	default:
		return false
	}
}

func (p *Pool) replace(c *client) {
	p.clients <- c
}

func (p *Pool) inc() bool {
	if p.created >= p.max {
		return false
	}

	p.mut.Lock()
	defer p.mut.Unlock()

	if p.created >= p.max {
		return false
	}
	p.created++
	return true
}

func (p *Pool) dec() {
	p.mut.Lock()
	p.created--
	p.mut.Unlock()

	select {
	case p.rebuild <- struct{}{}:
	default:
	}
}

func (p *Pool) makeOne() {
	go func() {
		if p.inc() {
			if c, err := p.build(); err == nil {
				p.clients <- c
			} else {
				p.lastBuildErr = &timestampedErr{err, time.Now()}
				p.dec()
			}
		}
	}()
}

func startTLS(c *client, addr string) (bool, error) {
	if ok, _ := c.Extension("STARTTLS"); !ok {
		return false, nil
	}

	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return false, err
	}

	if err := c.StartTLS(&tls.Config{ServerName: host}); err != nil {
		return false, err
	}

	return true, nil
}

func addAuth(c *client, auth smtp.Auth) (bool, error) {
	if ok, _ := c.Extension("AUTH"); !ok {
		return false, nil
	}

	if err := c.Auth(auth); err != nil {
		return false, err
	}

	return true, nil
}

func (p *Pool) build() (*client, error) {
	cl, err := smtp.Dial(p.addr)
	if err != nil {
		return nil, err
	}
	c := &client{cl, 0}

	if _, err := startTLS(c, p.addr); err != nil {
		c.Close()
		return nil, err
	}

	if p.auth != nil {
		if _, err := addAuth(c, p.auth); err != nil {
			c.Close()
			return nil, err
		}
	}

	return c, nil
}

func (p *Pool) maybeReplace(err error, c *client) {
	if err == nil {
		c.failCount = 0
		p.replace(c)
		return
	}

	c.failCount++
	if c.failCount >= maxFails {
		goto shutdown
	}

	if !shouldReuse(err) {
		goto shutdown
	}

	if err := c.Reset(); err != nil {
		goto shutdown
	}

	p.replace(c)
	return

shutdown:
	p.dec()
	c.Close()
}

func (p *Pool) failedToGet(startTime time.Time) error {
	select {
	case <-p.closing:
		return ErrClosed
	default:
	}

	if p.lastBuildErr != nil && startTime.Before(p.lastBuildErr.ts) {
		return p.lastBuildErr.err
	}

	return ErrTimeout
}

// Send sends an email via a connection pulled from the Pool. The timeout may
// be <0 to indicate no timeout. Otherwise reaching the timeout will produce
// and error building a connection that occurred while we were waiting, or
// otherwise ErrTimeout.
func (p *Pool) Send(e *Email, timeout time.Duration) (err error) {
	start := time.Now()
	c := p.get(timeout)
	if c == nil {
		return p.failedToGet(start)
	}

	defer func() {
		p.maybeReplace(err, c)
	}()

	recipients, err := addressLists(e.To, e.Cc, e.Bcc)
	if err != nil {
		return
	}

	msg, err := e.Bytes()
	if err != nil {
		return
	}

	from, err := emailOnly(e.From)
	if err != nil {
		return
	}
	if err = c.Mail(from); err != nil {
		return
	}

	for _, recip := range recipients {
		if err = c.Rcpt(recip); err != nil {
			return
		}
	}

	w, err := c.Data()
	if err != nil {
		return
	}
	if _, err = w.Write(msg); err != nil {
		return
	}

	err = w.Close()

	return
}

func emailOnly(full string) (string, error) {
	addr, err := mail.ParseAddress(full)
	if err != nil {
		return "", err
	}
	return addr.Address, nil
}

func addressLists(lists ...[]string) ([]string, error) {
	length := 0
	for _, lst := range lists {
		length += len(lst)
	}
	combined := make([]string, 0, length)

	for _, lst := range lists {
		for _, full := range lst {
			addr, err := emailOnly(full)
			if err != nil {
				return nil, err
			}
			combined = append(combined, addr)
		}
	}

	return combined, nil
}

// Close immediately changes the pool's state so no new connections will be
// created, then gets and closes the existing ones as they become available.
func (p *Pool) Close() {
	close(p.closing)

	for p.created > 0 {
		c := <-p.clients
		c.Quit()
		p.dec()
	}
}
<|endoftext|>


<commit_before>package procspy

// /proc-based implementation.

import (
	"bytes"
	"os"
	"strconv"
	"syscall"
)

const (
	procRoot = "/proc"
)

// walkProcPid walks over all numerical (PID) /proc entries, and sees if their
// ./fd/* files are symlink to sockets. Returns a map from socket ID (inode)
// to PID. Will return an error if /proc isn't there.
func walkProcPid() (map[uint64]Proc, error) {
	fh, err := os.Open(procRoot)
	if err != nil {
		return nil, err
	}

	dirNames, err := fh.Readdirnames(-1)
	fh.Close()
	if err != nil {
		return nil, err
	}

	var (
		res  = map[uint64]Proc{}
		stat syscall.Stat_t
	)
	for _, dirName := range dirNames {
		pid, err := strconv.ParseUint(dirName, 10, 0)
		if err != nil {
			// Not a number, so not a PID subdir.
			continue
		}

		fdBase := procRoot + "/" + dirName + "/fd/"
		dfh, err := os.Open(fdBase)
		if err != nil {
			// Process is be gone by now, or we don't have access.
			continue
		}

		fdNames, err := dfh.Readdirnames(-1)
		dfh.Close()
		if err != nil {
			continue
		}

		var name string

		for _, fdName := range fdNames {
			// Direct use of syscall.Stat() to save garbage.
			err = syscall.Stat(fdBase+fdName, &stat)
			if err != nil {
				continue
			}

			// We want sockets only.
			if stat.Mode&syscall.S_IFMT != syscall.S_IFSOCK {
				continue
			}

			if name == "" {
				if name = procName(procRoot + "/" + dirName); name == "" {
					// Process might be gone by now
					break
				}
			}

			res[stat.Ino] = Proc{
				PID:  uint(pid),
				Name: name,
			}
		}
	}

	return res, nil
}

// procName does a pid->name lookup.
func procName(base string) string {
	fh, err := os.Open(base + "/comm")
	if err != nil {
		return ""
	}

	name := make([]byte, 64)
	l, err := fh.Read(name)
	fh.Close()
	if err != nil {
		return ""
	}

	if l < 2 {
		return ""
	}

	// drop trailing "\n"
	return string(name[:l-1])
}

// readFile reads an arbitrary file into a buffer. It's a variable so it can
// be overwritten for benchmarks. That's bad practice and we should change it
// to be a dependency.
var readFile = func(filename string, buf *bytes.Buffer) error {
	f, err := os.Open(filename)
	if err != nil {
		return err
	}
	_, err = buf.ReadFrom(f)
	f.Close()
	return err
}
<commit_msg>Allow consumer to specify where the proc filesystem is.<commit_after>package procspy

// /proc-based implementation.

import (
	"bytes"
	"os"
	"strconv"
	"syscall"
)

var (
	procRoot = "/proc"
)

// SetProcRoot sets the location of the proc filesystem.
func SetProcRoot(root string) {
	procRoot = root
}

// walkProcPid walks over all numerical (PID) /proc entries, and sees if their
// ./fd/* files are symlink to sockets. Returns a map from socket ID (inode)
// to PID. Will return an error if /proc isn't there.
func walkProcPid() (map[uint64]Proc, error) {
	fh, err := os.Open(procRoot)
	if err != nil {
		return nil, err
	}

	dirNames, err := fh.Readdirnames(-1)
	fh.Close()
	if err != nil {
		return nil, err
	}

	var (
		res  = map[uint64]Proc{}
		stat syscall.Stat_t
	)
	for _, dirName := range dirNames {
		pid, err := strconv.ParseUint(dirName, 10, 0)
		if err != nil {
			// Not a number, so not a PID subdir.
			continue
		}

		fdBase := procRoot + "/" + dirName + "/fd/"
		dfh, err := os.Open(fdBase)
		if err != nil {
			// Process is be gone by now, or we don't have access.
			continue
		}

		fdNames, err := dfh.Readdirnames(-1)
		dfh.Close()
		if err != nil {
			continue
		}

		var name string

		for _, fdName := range fdNames {
			// Direct use of syscall.Stat() to save garbage.
			err = syscall.Stat(fdBase+fdName, &stat)
			if err != nil {
				continue
			}

			// We want sockets only.
			if stat.Mode&syscall.S_IFMT != syscall.S_IFSOCK {
				continue
			}

			if name == "" {
				if name = procName(procRoot + "/" + dirName); name == "" {
					// Process might be gone by now
					break
				}
			}

			res[stat.Ino] = Proc{
				PID:  uint(pid),
				Name: name,
			}
		}
	}

	return res, nil
}

// procName does a pid->name lookup.
func procName(base string) string {
	fh, err := os.Open(base + "/comm")
	if err != nil {
		return ""
	}

	name := make([]byte, 64)
	l, err := fh.Read(name)
	fh.Close()
	if err != nil {
		return ""
	}

	if l < 2 {
		return ""
	}

	// drop trailing "\n"
	return string(name[:l-1])
}

// readFile reads an arbitrary file into a buffer. It's a variable so it can
// be overwritten for benchmarks. That's bad practice and we should change it
// to be a dependency.
var readFile = func(filename string, buf *bytes.Buffer) error {
	f, err := os.Open(filename)
	if err != nil {
		return err
	}
	_, err = buf.ReadFrom(f)
	f.Close()
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"
	"sync"

	"gopkg.in/itchio/rsync-go.v0"

	"github.com/itchio/butler/bio"
	"github.com/itchio/butler/wharf"
)

func push(src string, repoSpec string) {
	err := doPush(src, repoSpec)
	if err != nil {
		bio.Die(err.Error())
	}
}

var (
	addressRe   = regexp.MustCompile("[^:]:.*")
	defaultPort = 22
)

func doPush(src string, repoSpec string) error {
	address := *pushAddress
	if !addressRe.MatchString(address) {
		address = fmt.Sprintf("%s:%d", address, defaultPort)
	}

	conn, err := wharf.Connect(address, *pushIdentity, version)
	if err != nil {
		return err
	}
	defer conn.Close()

	up := bio.UploadParams{RepoSpec: repoSpec}
	ok, _, err := conn.SendRequest("butler/upload-params", true, up)
	if err != nil {
		return fmt.Errorf("Could not send upload params: %s", err.Error())
	}

	if !ok {
		return fmt.Errorf("Could not find upload to replace from '%s'", repoSpec)
	}
	bio.Log("upload params were accepted!")

	var wg sync.WaitGroup

	done := make(chan bool)
	errs := make(chan error)

	for j := 0; j < 10; j++ {
		wg.Add(1)
		j := j
		go func() {
			defer wg.Done()
			path := fmt.Sprintf("/test/channel/%d", j)
			ch, err := conn.OpenCompressedChannel("butler/send-file", bio.FileAdded{Path: path})
			if err != nil {
				errs <- err
				return
			}
			defer ch.Close()

			for i := 0; i < 250; i++ {
				err := ch.Send(bio.SourceFile{
					Path:   fmt.Sprintf(""),
					Hashes: []rsync.BlockHash{},
				})
				if err != nil {
					errs <- err
					return
				}
			}
		}()
	}

	go func() {
		wg.Wait()
		done <- true
	}()

	select {
	case err := <-errs:
		return err
	case <-done:
		return nil
	}
}
<commit_msg>Walk dir<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"

	"gopkg.in/itchio/rsync-go.v0"

	"github.com/dustin/go-humanize"
	"github.com/itchio/butler/bio"
	"github.com/itchio/butler/wharf"
)

func push(src string, repoSpec string) {
	err := doPush(src, repoSpec)
	if err != nil {
		bio.Die(err.Error())
	}
}

var (
	addressRe   = regexp.MustCompile("[^:]:.*")
	defaultPort = 22
)

func doPush(src string, repoSpec string) error {
	address := *pushAddress
	if !addressRe.MatchString(address) {
		address = fmt.Sprintf("%s:%d", address, defaultPort)
	}

	conn, err := wharf.Connect(address, *pushIdentity, version)
	if err != nil {
		return err
	}
	defer conn.Close()

	up := bio.UploadParams{RepoSpec: repoSpec}
	ok, _, err := conn.SendRequest("butler/upload-params", true, up)
	if err != nil {
		return fmt.Errorf("Could not send upload params: %s", err.Error())
	}

	if !ok {
		return fmt.Errorf("Could not find upload to replace from '%s'", repoSpec)
	}
	bio.Log("upload params were accepted!")

	var wg sync.WaitGroup

	done := make(chan bool)
	errs := make(chan error)

	log.Printf("\n")
	var totalSize int64
	fileList := make(map[string]os.FileInfo)

	filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			errs <- err
			return err
		}

		if info.IsDir() && len(path) > 1 && strings.HasPrefix(info.Name(), ".") {
			log.Printf("Skipping %s because it starts with .\n", path)
			return filepath.SkipDir
		}

		if !info.IsDir() {
			fileList[path] = info
		}
		return nil
	})

	for path, info := range fileList {
		totalSize += info.Size()
		humanSize := humanize.Bytes(uint64(info.Size()))
		log.Printf("%8s | %s\n", humanSize, path)
	}
	log.Printf("---------+-------------\n")
	log.Printf("%8s | (total)\n", humanize.Bytes(uint64(totalSize)))

	for j := 0; j < 10; j++ {
		wg.Add(1)
		j := j
		go func() {
			defer wg.Done()
			path := fmt.Sprintf("/test/channel/%d", j)
			ch, err := conn.OpenCompressedChannel("butler/send-file", bio.FileAdded{Path: path})
			if err != nil {
				errs <- err
				return
			}
			defer ch.Close()

			for i := 0; i < 250; i++ {
				err := ch.Send(bio.SourceFile{
					Path:   fmt.Sprintf(""),
					Hashes: []rsync.BlockHash{},
				})
				if err != nil {
					errs <- err
					return
				}
			}
		}()
	}

	go func() {
		wg.Wait()
		done <- true
	}()

	select {
	case err := <-errs:
		return err
	case <-done:
		return nil
	}
}
<|endoftext|>


<commit_before>package proxmox

import (
	"errors"
	"net/url"
	"strconv"
	"strings"
	"time"

	_ "github.com/davecgh/go-spew/spew"
)

type QemuVM struct {
	Mem       float64 `json:"mem"`
	CPUs      float64 `json:"cpus"`
	NetOut    float64 `json:"netout"`
	PID       string  `json:"pid"`
	Disk      float64 `json:"disk"`
	MaxMem    float64 `json:"maxmem"`
	Status    string  `json:"status"`
	Template  float64 `json:"template"`
	NetIn     float64 `json:"netin"`
	MaxDisk   float64 `json:"maxdisk"`
	Name      string  `json:"name"`
	DiskWrite float64 `json:"diskwrite"`
	CPU       float64 `json:"cpu"`
	VMId      float64 `json:"vmid"`
	DiskRead  float64 `json:"diskread"`
	Uptime    float64 `json:"uptime"`
	Node      Node
}

type QemuList map[string]QemuVM

type QemuNet map[string]string

type QemuConfig struct {
	Bootdisk string  `json:"bootdisk"`
	Cores    float64 `json:"cores"`
	Digest   string  `json:"digest"`
	Memory   float64 `json:"memory"`
	Net      map[string]QemuNet
	SMBios1  string            `json:"smbios1"`
	Sockets  float64           `json:"sockets"`
	Disks    map[string]string `json:"disks"`
}

type QemuStatus struct {
	CPU       float64 `json:"cpu"`
	CPUs      float64 `json:"cpus"`
	Mem       float64 `json:"mem"`
	MaxMem    float64 `json:"maxmem"`
	Disk      float64 `json:"disk"`
	MaxDisk   float64 `json:"maxdisk"`
	DiskWrite float64 `json:"diskwrite"`
	DiskRead  float64 `json:"diskread"`
	NetIn     float64 `json:"netin"`
	NetOut    float64 `json:"netout"`
	Uptime    float64 `json:"uptime"`
	QmpStatus string  `json:"qmpstatus"`
	Status    string  `json:"status"`
	Template  string  `json:"template"`
}

func (qemu QemuVM) Delete() (map[string]interface{}, error) {
	var target string
	var data map[string]interface{}
	var err error

	//fmt.Print("!QemuDelete ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64)
	data, err = qemu.Node.Proxmox.Delete(target)
	if err != nil {
		return nil, err
	}
	return data, nil
}

func stringToMap(data string, itemSeparator string, kvSeparator string) map[string]string {
	var result map[string]string

	result = make(map[string]string)
	list := strings.Split(data, itemSeparator)
	for _, item := range list {
		kv := strings.Split(item, kvSeparator)
		result[kv[0]] = kv[1]
	}
	return result
}

func (qemu QemuVM) Config() (QemuConfig, error) {
	var target string
	var data map[string]interface{}
	var results map[string]interface{}
	var config QemuConfig
	var err error

	//fmt.Print("!QemuConfig ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/config"
	data, err = qemu.Node.Proxmox.Get(target)
	results = data["data"].(map[string]interface{})
	if err != nil {
		return config, err
	}
	config = QemuConfig{
		Bootdisk: results["bootdisk"].(string),
		Cores:    results["cores"].(float64),
		Digest:   results["digest"].(string),
		Memory:   results["memory"].(float64),
		Sockets:  results["sockets"].(float64),
		SMBios1:  results["smbios1"].(string),
	}
	disktype := [3]string{"virtio", "sata", "ide"}
	disknum := [4]string{"0", "1", "2", "3"}
	config.Disks = make(map[string]string)
	for _, d := range disktype {
		for _, i := range disknum {
			id := d + i
			if disk, ok := results[id]; ok {
				config.Disks[id] = disk.(string)
			}
		}
	}
	config.Net = make(map[string]QemuNet)
	netnum := [4]string{"0", "1", "2", "3"}
	for _, n := range netnum {
		if net, ok := results["net"+n]; ok {
			config.Net["net"+n] = stringToMap(net.(string), ",", "=")
		}
	}

	return config, nil
}

func (qemu QemuVM) CurrentStatus() (QemuStatus, error) {
	var target string
	var err error
	var data map[string]interface{}
	var results map[string]interface{}
	var status QemuStatus

	//fmt.Println("!QemuStatus ", strconv.FormatFloat(qemu.VMId, 'f', 0, 64))

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/current"
	data, err = qemu.Node.Proxmox.Get(target)
	if err != nil {
		return status, err
	}
	results = data["data"].(map[string]interface{})
	status = QemuStatus{
		CPU:       results["cpu"].(float64),
		CPUs:      results["cpus"].(float64),
		Mem:       results["mem"].(float64),
		MaxMem:    results["maxmem"].(float64),
		Disk:      results["disk"].(float64),
		MaxDisk:   results["maxdisk"].(float64),
		DiskWrite: results["diskwrite"].(float64),
		DiskRead:  results["diskread"].(float64),
		NetIn:     results["netin"].(float64),
		NetOut:    results["netout"].(float64),
		Uptime:    results["uptime"].(float64),
		QmpStatus: results["qmpstatus"].(string),
		Status:    results["status"].(string),
		Template:  results["template"].(string),
	}
	return status, nil
}

func (qemu QemuVM) WaitForStatus(status string, timeout int) error {
	var i int
	var err error
	var qStatus QemuStatus
	for i = 0; i < timeout; i++ {
		qStatus, err = qemu.CurrentStatus()
		if err != nil {
			return err
		}

		if qStatus.Status == status {
			return nil
		}
		time.Sleep(time.Second * 1)
	}
	return errors.New("Timeout reached")
}

func (qemu QemuVM) Start() error {
	var target string
	var err error

	//fmt.Println("!QemuStart ", strconv.FormatFloat(qemu.VMId, 'f', 0, 64))

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/start"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Stop() error {
	var target string
	var err error

	//fmt.Print("!QemuStop ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/stop"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Shutdown() error {
	var target string
	var err error

	//fmt.Print("!QemuShutdown ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/shutdown"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Suspend() error {
	var target string
	var err error

	//fmt.Print("!QemuSuspend ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/suspend"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Resume() error {
	var target string
	var err error

	//fmt.Print("!QemuResume ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/resume"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Clone(newId float64, name string, targetName string) (string, error) {
	var target string
	var err error

	newVMID := strconv.FormatFloat(newId, 'f', 0, 64)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/clone"

	form := url.Values{
		"newid":  {newVMID},
		"name":   {name},
		"target": {targetName},
		"full":   {"1"},
	}

	data, err := qemu.Node.Proxmox.PostForm(target, form)
	if err != err {
		return "", err
	}

	UPid := data["data"].(string)

	return UPid, nil
}
<commit_msg>return Task instead of string when cloning<commit_after>package proxmox

import (
	"errors"
	"net/url"
	"strconv"
	"strings"
	"time"

	_ "github.com/davecgh/go-spew/spew"
)

type QemuVM struct {
	Mem       float64 `json:"mem"`
	CPUs      float64 `json:"cpus"`
	NetOut    float64 `json:"netout"`
	PID       string  `json:"pid"`
	Disk      float64 `json:"disk"`
	MaxMem    float64 `json:"maxmem"`
	Status    string  `json:"status"`
	Template  float64 `json:"template"`
	NetIn     float64 `json:"netin"`
	MaxDisk   float64 `json:"maxdisk"`
	Name      string  `json:"name"`
	DiskWrite float64 `json:"diskwrite"`
	CPU       float64 `json:"cpu"`
	VMId      float64 `json:"vmid"`
	DiskRead  float64 `json:"diskread"`
	Uptime    float64 `json:"uptime"`
	Node      Node
}

type QemuList map[string]QemuVM

type QemuNet map[string]string

type QemuConfig struct {
	Bootdisk string  `json:"bootdisk"`
	Cores    float64 `json:"cores"`
	Digest   string  `json:"digest"`
	Memory   float64 `json:"memory"`
	Net      map[string]QemuNet
	SMBios1  string            `json:"smbios1"`
	Sockets  float64           `json:"sockets"`
	Disks    map[string]string `json:"disks"`
}

type QemuStatus struct {
	CPU       float64 `json:"cpu"`
	CPUs      float64 `json:"cpus"`
	Mem       float64 `json:"mem"`
	MaxMem    float64 `json:"maxmem"`
	Disk      float64 `json:"disk"`
	MaxDisk   float64 `json:"maxdisk"`
	DiskWrite float64 `json:"diskwrite"`
	DiskRead  float64 `json:"diskread"`
	NetIn     float64 `json:"netin"`
	NetOut    float64 `json:"netout"`
	Uptime    float64 `json:"uptime"`
	QmpStatus string  `json:"qmpstatus"`
	Status    string  `json:"status"`
	Template  string  `json:"template"`
}

func (qemu QemuVM) Delete() (map[string]interface{}, error) {
	var target string
	var data map[string]interface{}
	var err error

	//fmt.Print("!QemuDelete ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64)
	data, err = qemu.Node.Proxmox.Delete(target)
	if err != nil {
		return nil, err
	}
	return data, nil
}

func stringToMap(data string, itemSeparator string, kvSeparator string) map[string]string {
	var result map[string]string

	result = make(map[string]string)
	list := strings.Split(data, itemSeparator)
	for _, item := range list {
		kv := strings.Split(item, kvSeparator)
		result[kv[0]] = kv[1]
	}
	return result
}

func (qemu QemuVM) Config() (QemuConfig, error) {
	var target string
	var data map[string]interface{}
	var results map[string]interface{}
	var config QemuConfig
	var err error

	//fmt.Print("!QemuConfig ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/config"
	data, err = qemu.Node.Proxmox.Get(target)
	results = data["data"].(map[string]interface{})
	if err != nil {
		return config, err
	}
	config = QemuConfig{
		Bootdisk: results["bootdisk"].(string),
		Cores:    results["cores"].(float64),
		Digest:   results["digest"].(string),
		Memory:   results["memory"].(float64),
		Sockets:  results["sockets"].(float64),
		SMBios1:  results["smbios1"].(string),
	}
	disktype := [3]string{"virtio", "sata", "ide"}
	disknum := [4]string{"0", "1", "2", "3"}
	config.Disks = make(map[string]string)
	for _, d := range disktype {
		for _, i := range disknum {
			id := d + i
			if disk, ok := results[id]; ok {
				config.Disks[id] = disk.(string)
			}
		}
	}
	config.Net = make(map[string]QemuNet)
	netnum := [4]string{"0", "1", "2", "3"}
	for _, n := range netnum {
		if net, ok := results["net"+n]; ok {
			config.Net["net"+n] = stringToMap(net.(string), ",", "=")
		}
	}

	return config, nil
}

func (qemu QemuVM) CurrentStatus() (QemuStatus, error) {
	var target string
	var err error
	var data map[string]interface{}
	var results map[string]interface{}
	var status QemuStatus

	//fmt.Println("!QemuStatus ", strconv.FormatFloat(qemu.VMId, 'f', 0, 64))

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/current"
	data, err = qemu.Node.Proxmox.Get(target)
	if err != nil {
		return status, err
	}
	results = data["data"].(map[string]interface{})
	status = QemuStatus{
		CPU:       results["cpu"].(float64),
		CPUs:      results["cpus"].(float64),
		Mem:       results["mem"].(float64),
		MaxMem:    results["maxmem"].(float64),
		Disk:      results["disk"].(float64),
		MaxDisk:   results["maxdisk"].(float64),
		DiskWrite: results["diskwrite"].(float64),
		DiskRead:  results["diskread"].(float64),
		NetIn:     results["netin"].(float64),
		NetOut:    results["netout"].(float64),
		Uptime:    results["uptime"].(float64),
		QmpStatus: results["qmpstatus"].(string),
		Status:    results["status"].(string),
		Template:  results["template"].(string),
	}
	return status, nil
}

func (qemu QemuVM) WaitForStatus(status string, timeout int) error {
	var i int
	var err error
	var qStatus QemuStatus
	for i = 0; i < timeout; i++ {
		qStatus, err = qemu.CurrentStatus()
		if err != nil {
			return err
		}

		if qStatus.Status == status {
			return nil
		}
		time.Sleep(time.Second * 1)
	}
	return errors.New("Timeout reached")
}

func (qemu QemuVM) Start() error {
	var target string
	var err error

	//fmt.Println("!QemuStart ", strconv.FormatFloat(qemu.VMId, 'f', 0, 64))

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/start"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Stop() error {
	var target string
	var err error

	//fmt.Print("!QemuStop ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/stop"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Shutdown() error {
	var target string
	var err error

	//fmt.Print("!QemuShutdown ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/shutdown"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Suspend() error {
	var target string
	var err error

	//fmt.Print("!QemuSuspend ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/suspend"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Resume() error {
	var target string
	var err error

	//fmt.Print("!QemuResume ", qemu.VMId)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/status/resume"
	_, err = qemu.Node.Proxmox.Post(target, "")
	return err
}

func (qemu QemuVM) Clone(newId float64, name string, targetName string) (Task, error) {
	var target string
	var err error

	newVMID := strconv.FormatFloat(newId, 'f', 0, 64)

	target = "nodes/" + qemu.Node.Node + "/qemu/" + strconv.FormatFloat(qemu.VMId, 'f', 0, 64) + "/clone"

	form := url.Values{
		"newid":  {newVMID},
		"name":   {name},
		"target": {targetName},
		"full":   {"1"},
	}

	data, err := qemu.Node.Proxmox.PostForm(target, form)
	if err != err {
		return Task{}, err
	}

	t := Task{
		UPid:    data["data"].(string),
		proxmox: qemu.Node.Proxmox,
	}

	return t, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/OWASP/Amass/v3/stringset"
)

const (
	// UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_2) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/72.0.3626.119 Safari/537.36"

	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	defaultTLSConnectTimeout = 3 * time.Second
	defaultHandshakeDeadline = 5 * time.Second
)

var (
	defaultClient *http.Client
)

func init() {
	jar, _ := cookiejar.New(nil)
	defaultClient = &http.Client{
		Timeout: 30 * time.Second,
		Transport: &http.Transport{
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
				DualStack: true,
			}).DialContext,
			MaxIdleConns:          200,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   20 * time.Second,
			ExpectContinueTimeout: 20 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	defaultClient.Jar.SetCookies(destURL, defaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range defaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for
// the urlstring parameter when successful.
func RequestWebPage(urlstring string, body io.Reader, hvals map[string]string, uid, secret string) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequest(method, urlstring, body)
	if err != nil {
		return "", err
	}
	if uid != "" && secret != "" {
		req.SetBasicAuth(uid, secret)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)
	if hvals != nil {
		for k, v := range hvals {
			req.Header.Set(k, v)
		}
	}

	resp, err := defaultClient.Do(req)
	if err != nil {
		return "", err
	} else if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		return "", errors.New(resp.Status)
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	return string(in), err
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		cfg := &tls.Config{InsecureSkipVerify: true}
		// Set the maximum time allowed for making the connection
		ctx, cancel := context.WithTimeout(context.Background(), defaultTLSConnectTimeout)
		defer cancel()
		// Obtain the connection
		d := net.Dialer{}
		conn, err := d.DialContext(ctx, "tcp", addr+":"+strconv.Itoa(port))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, cfg)
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		// This goroutine will break us out of the handshake
		time.AfterFunc(defaultHandshakeDeadline, func() {
			errChan <- errors.New("Handshake timeout")
		})
		// Be sure we do not wait too long in this attempt
		c.SetDeadline(time.Now().Add(defaultHandshakeDeadline))
		// The handshake is performed in the goroutine
		go func() {
			errChan <- c.Handshake()
		}()
		// The error channel returns handshake or timeout error
		if err = <-errChan; err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}
	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode() string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage("https://ipapi.co/json", nil, headers, "", "")
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	json.Unmarshal([]byte(page), &ipinfo)
	return strings.ToLower(ipinfo.CountryCode)
}
<commit_msg>update to the HTTP client implementation<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/OWASP/Amass/v3/stringset"
)

const (
	// UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.97 Safari/537.36"

	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	defaultTLSConnectTimeout = 3 * time.Second
	defaultHandshakeDeadline = 5 * time.Second
)

var (
	defaultClient *http.Client
)

func init() {
	jar, _ := cookiejar.New(nil)
	defaultClient = &http.Client{
		Timeout: 30 * time.Second,
		Transport: &http.Transport{
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
				DualStack: true,
			}).DialContext,
			MaxIdleConns:          200,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   20 * time.Second,
			ExpectContinueTimeout: 20 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	defaultClient.Jar.SetCookies(destURL, defaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range defaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for
// the urlstring parameter when successful.
func RequestWebPage(urlstring string, body io.Reader, hvals map[string]string, uid, secret string) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequest(method, urlstring, body)
	if err != nil {
		return "", err
	}
	if uid != "" && secret != "" {
		req.SetBasicAuth(uid, secret)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)
	if hvals != nil {
		for k, v := range hvals {
			req.Header.Set(k, v)
		}
	}

	resp, err := defaultClient.Do(req)
	if err != nil {
		return "", err
	} else if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		return "", errors.New(resp.Status)
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	return string(in), err
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		cfg := &tls.Config{InsecureSkipVerify: true}
		// Set the maximum time allowed for making the connection
		ctx, cancel := context.WithTimeout(context.Background(), defaultTLSConnectTimeout)
		defer cancel()
		// Obtain the connection
		d := net.Dialer{}
		conn, err := d.DialContext(ctx, "tcp", addr+":"+strconv.Itoa(port))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, cfg)
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		// This goroutine will break us out of the handshake
		time.AfterFunc(defaultHandshakeDeadline, func() {
			errChan <- errors.New("Handshake timeout")
		})
		// Be sure we do not wait too long in this attempt
		c.SetDeadline(time.Now().Add(defaultHandshakeDeadline))
		// The handshake is performed in the goroutine
		go func() {
			errChan <- c.Handshake()
		}()
		// The error channel returns handshake or timeout error
		if err = <-errChan; err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}
	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode() string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage("https://ipapi.co/json", nil, headers, "", "")
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	json.Unmarshal([]byte(page), &ipinfo)
	return strings.ToLower(ipinfo.CountryCode)
}
<|endoftext|>