<commit_before>/*
 * Copyright 2014 Jason Woods.
 *
 * This file is a modification of code from Logstash Forwarder.
 * Copyright 2012-2013 Jordan Sissel and contributors.
 *
 * 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 main

import (
  "bufio"
  "bytes"
  "io"
  "log"
  "os"
  "time"
)

type Harvester struct {
  info       *ProspectorInfo
  fileinfo   os.FileInfo
  path       string /* the file path to harvest */
  fileconfig *FileConfig
  offset     int64
  codec      Codec

  file *os.File /* the file being watched */
}

func NewHarvester(info *ProspectorInfo, fileconfig *FileConfig, offset int64) *Harvester {
  var fileinfo os.FileInfo
  var path string
  if info != nil {
    // Grab now so we can safely use them even if prospector changes them
    fileinfo = info.identity.Stat()
    path = info.file
  } else {
    // This is stdin
    fileinfo = nil
    path = "-"
  }
  return &Harvester{
    info:       info,
    fileinfo:   fileinfo,
    path:       path,
    fileconfig: fileconfig,
    offset:     offset,
  }
}

func (h *Harvester) Harvest(output chan<- *FileEvent) (int64, bool) {
  if !h.prepareHarvester() {
    return h.offset, true
  }

  h.codec = h.fileconfig.codec.NewCodec(h.path, h.fileconfig, h.info, h.offset, output)

  defer h.file.Close()

  // NOTE(driskell): How would we know line number if from_beginning is false and we SEEK_END? Or would we scan,count,skip?
  var line uint64 = 0 // Ask registrar about the line number

  // Get current offset in file
  // TODO: Check error?
  offset, _ := h.file.Seek(0, os.SEEK_CUR)
  log.Printf("Started harvester at position %d (requested %d): %s\n", offset, h.offset, h.path)
  h.offset = offset

  // TODO(sissel): Make the buffer size tunable at start-time
  reader := bufio.NewReaderSize(h.file, 16<<10) // 16kb buffer by default
  buffer := new(bytes.Buffer)

  var read_timeout = 10 * time.Second
  last_read_time := time.Now()

ReadLoop:
  for {
    // Check shutdown
    select {
    case <-h.info.ShutdownSignal():
      break ReadLoop
    default:
    }

    text, bytesread, err := h.readline(reader, buffer, read_timeout)

    if err != nil {
      if err == io.EOF {
        // Timed out waiting for data, got EOF, check to see if the file was truncated
        if h.path == "-" {
          // This wouldn't make sense on stdin so lets not risk anything strange happening
          continue
        }

        info, err := h.file.Stat()
        if err == nil {
          if info.Size() < h.offset {
            log.Printf("File truncated, seeking to beginning: %s\n", h.path)
            h.file.Seek(0, os.SEEK_SET)
            h.offset = 0
            continue
          } else if age := time.Since(last_read_time); age > h.fileconfig.DeadTime {
            // if last_read_time was more than dead time, this file is probably dead. Stop watching it.
            log.Printf("Stopping harvest of %s; last change was %v ago\n", h.path, age-(age%time.Second))
            return h.codec.Teardown(), false
          }

          continue
        } else {
          log.Printf("Unexpected error checking status of %s: %s\n", h.path, err)
        }
      } else {
        log.Printf("Unexpected error reading from %s: %s\n", h.path, err)
      }
      return h.codec.Teardown(), true
    }
    last_read_time = time.Now()

    line++
    line_offset := h.offset
    h.offset += int64(bytesread)

    // Codec is last - it saves harvester state for us such as offset for resume
    h.codec.Event(line_offset, h.offset, line, text)
  }

  log.Printf("Harvester shutdown for %s complete\n", h.path)
  return h.codec.Teardown(), false
}

func (h *Harvester) prepareHarvester() bool {
  // Special handling that "-" means to read from standard input
  if h.path == "-" {
    h.file = os.Stdin
    return true
  }

  var err error
  h.file, err = h.openFile(h.path)
  if err != nil {
    log.Printf("Failed opening %s: %s\n", h.path, err)
    return false
  }

  // Check we opened the right file
  info, err := h.file.Stat()
  if err != nil || !os.SameFile(info, h.fileinfo) {
    h.file.Close()
    return false
  }

  // TODO: Check error?
  h.file.Seek(h.offset, os.SEEK_SET)

  return true
}

func (h *Harvester) readline(reader *bufio.Reader, buffer *bytes.Buffer, eof_timeout time.Duration) (*string, int, error) {
  var is_partial bool = true
  var newline_length int = 1
  start_time := time.Now()

  for {
    segment, err := reader.ReadBytes('\n')

    if segment != nil && len(segment) > 0 {
      if segment[len(segment)-1] == '\n' {
        // Found a complete line
        is_partial = false

        // Check if also a CR present
        if len(segment) > 1 && segment[len(segment)-2] == '\r' {
          newline_length++
        }
      }

      // TODO(sissel): if buffer exceeds a certain length, maybe report an error condition? chop it?
      buffer.Write(segment)
    }

    if err != nil {
      if err == io.EOF && is_partial {
        // Backoff
        select {
        case <-h.info.ShutdownSignal():
          return nil, 0, err
        case <-time.After(1 * time.Second):
        }

        // Give up waiting for data after a certain amount of time.
        // If we time out, return the error (eof)
        if time.Since(start_time) > eof_timeout {
          return nil, 0, err
        }
        continue
      } else {
        log.Println(err)
        return nil, 0, err // TODO(sissel): don't do this?
      }
    }

    // If we got a full line, return the whole line without the EOL chars (CRLF or LF)
    if !is_partial {
      // Get the str length with the EOL chars (LF or CRLF)
      bufferSize := buffer.Len()
      str := new(string)
      *str = buffer.String()[:bufferSize-newline_length]
      // Reset the buffer for the next line
      buffer.Reset()
      return str, bufferSize, nil
    }
  } /* forever read chunks */

  return nil, 0, nil
}
<commit_msg>Note a possible race condition<commit_after>/*
 * Copyright 2014 Jason Woods.
 *
 * This file is a modification of code from Logstash Forwarder.
 * Copyright 2012-2013 Jordan Sissel and contributors.
 *
 * 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 main

import (
  "bufio"
  "bytes"
  "io"
  "log"
  "os"
  "time"
)

type Harvester struct {
  info       *ProspectorInfo
  fileinfo   os.FileInfo
  path       string /* the file path to harvest */
  fileconfig *FileConfig
  offset     int64
  codec      Codec

  file *os.File /* the file being watched */
}

func NewHarvester(info *ProspectorInfo, fileconfig *FileConfig, offset int64) *Harvester {
  var fileinfo os.FileInfo
  var path string
  if info != nil {
    // Grab now so we can safely use them even if prospector changes them
    fileinfo = info.identity.Stat()
    path = info.file
  } else {
    // This is stdin
    fileinfo = nil
    path = "-"
  }
  return &Harvester{
    info:       info,
    fileinfo:   fileinfo,
    path:       path,
    fileconfig: fileconfig,
    offset:     offset,
  }
}

func (h *Harvester) Harvest(output chan<- *FileEvent) (int64, bool) {
  if !h.prepareHarvester() {
    return h.offset, true
  }

  h.codec = h.fileconfig.codec.NewCodec(h.path, h.fileconfig, h.info, h.offset, output)

  defer h.file.Close()

  // NOTE(driskell): How would we know line number if from_beginning is false and we SEEK_END? Or would we scan,count,skip?
  var line uint64 = 0 // Ask registrar about the line number

  // Get current offset in file
  // TODO: Check error?
  offset, _ := h.file.Seek(0, os.SEEK_CUR)
  log.Printf("Started harvester at position %d (requested %d): %s\n", offset, h.offset, h.path)
  h.offset = offset

  // TODO(sissel): Make the buffer size tunable at start-time
  reader := bufio.NewReaderSize(h.file, 16<<10) // 16kb buffer by default
  buffer := new(bytes.Buffer)

  var read_timeout = 10 * time.Second
  last_read_time := time.Now()

ReadLoop:
  for {
    // Check shutdown
    select {
    case <-h.info.ShutdownSignal():
      break ReadLoop
    default:
    }

    text, bytesread, err := h.readline(reader, buffer, read_timeout)

    if err != nil {
      if err == io.EOF {
        // Timed out waiting for data, got EOF, check to see if the file was truncated
        if h.path == "-" {
          // This wouldn't make sense on stdin so lets not risk anything strange happening
          continue
        }

        info, err := h.file.Stat()
        if err == nil {
          if info.Size() < h.offset {
            log.Printf("File truncated, seeking to beginning: %s\n", h.path)
            h.file.Seek(0, os.SEEK_SET)
            h.offset = 0
            continue
          } else if age := time.Since(last_read_time); age > h.fileconfig.DeadTime {
            // if last_read_time was more than dead time, this file is probably dead. Stop watching it.
            log.Printf("Stopping harvest of %s; last change was %v ago\n", h.path, age-(age%time.Second))
            // TODO: We should return a Stat() from before we attempted to read
            // In prospector we use that for comparison to resume
            // This prevents a potential race condition if we stop just as the
            // file is modified with extra lines...
            return h.codec.Teardown(), false
          }

          continue
        } else {
          log.Printf("Unexpected error checking status of %s: %s\n", h.path, err)
        }
      } else {
        log.Printf("Unexpected error reading from %s: %s\n", h.path, err)
      }
      return h.codec.Teardown(), true
    }
    last_read_time = time.Now()

    line++
    line_offset := h.offset
    h.offset += int64(bytesread)

    // Codec is last - it saves harvester state for us such as offset for resume
    h.codec.Event(line_offset, h.offset, line, text)
  }

  log.Printf("Harvester shutdown for %s complete\n", h.path)
  return h.codec.Teardown(), false
}

func (h *Harvester) prepareHarvester() bool {
  // Special handling that "-" means to read from standard input
  if h.path == "-" {
    h.file = os.Stdin
    return true
  }

  var err error
  h.file, err = h.openFile(h.path)
  if err != nil {
    log.Printf("Failed opening %s: %s\n", h.path, err)
    return false
  }

  // Check we opened the right file
  info, err := h.file.Stat()
  if err != nil || !os.SameFile(info, h.fileinfo) {
    h.file.Close()
    return false
  }

  // TODO: Check error?
  h.file.Seek(h.offset, os.SEEK_SET)

  return true
}

func (h *Harvester) readline(reader *bufio.Reader, buffer *bytes.Buffer, eof_timeout time.Duration) (*string, int, error) {
  var is_partial bool = true
  var newline_length int = 1
  start_time := time.Now()

  for {
    segment, err := reader.ReadBytes('\n')

    if segment != nil && len(segment) > 0 {
      if segment[len(segment)-1] == '\n' {
        // Found a complete line
        is_partial = false

        // Check if also a CR present
        if len(segment) > 1 && segment[len(segment)-2] == '\r' {
          newline_length++
        }
      }

      // TODO(sissel): if buffer exceeds a certain length, maybe report an error condition? chop it?
      buffer.Write(segment)
    }

    if err != nil {
      if err == io.EOF && is_partial {
        // Backoff
        select {
        case <-h.info.ShutdownSignal():
          return nil, 0, err
        case <-time.After(1 * time.Second):
        }

        // Give up waiting for data after a certain amount of time.
        // If we time out, return the error (eof)
        if time.Since(start_time) > eof_timeout {
          return nil, 0, err
        }
        continue
      } else {
        log.Println(err)
        return nil, 0, err // TODO(sissel): don't do this?
      }
    }

    // If we got a full line, return the whole line without the EOL chars (CRLF or LF)
    if !is_partial {
      // Get the str length with the EOL chars (LF or CRLF)
      bufferSize := buffer.Len()
      str := new(string)
      *str = buffer.String()[:bufferSize-newline_length]
      // Reset the buffer for the next line
      buffer.Reset()
      return str, bufferSize, nil
    }
  } /* forever read chunks */

  return nil, 0, nil
}
<|endoftext|>


<commit_before>package vault

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/vault/helper/forwarding"
	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/keepalive"
)

const (
	clusterListenerAcceptDeadline = 500 * time.Millisecond
	heartbeatInterval             = 30 * time.Second
)

// Starts the listeners and servers necessary to handle forwarded requests
func (c *Core) startForwarding() error {
	c.logger.Trace("core: cluster listener setup function")
	defer c.logger.Trace("core: leaving cluster listener setup function")

	// Clean up in case we have transitioned from a client to a server
	c.requestForwardingConnectionLock.Lock()
	c.clearForwardingClients()
	c.requestForwardingConnectionLock.Unlock()

	// Resolve locally to avoid races
	ha := c.ha != nil

	// Get our TLS config
	tlsConfig, err := c.ClusterTLSConfig()
	if err != nil {
		c.logger.Error("core: failed to get tls configuration when starting forwarding", "error", err)
		return err
	}

	// The server supports all of the possible protos
	tlsConfig.NextProtos = []string{"h2", "req_fw_sb-act_v1"}

	// Create our RPC server and register the request handler server
	c.clusterParamsLock.Lock()

	if c.rpcServer != nil {
		c.logger.Warn("core: forwarding rpc server already running")
		return nil
	}

	c.rpcServer = grpc.NewServer(
		grpc.KeepaliveParams(keepalive.ServerParameters{
			Time: 2 * heartbeatInterval,
		}),
	)

	if ha && c.clusterHandler != nil {
		RegisterRequestForwardingServer(c.rpcServer, &forwardedRequestRPCServer{
			core:    c,
			handler: c.clusterHandler,
		})
	}
	c.clusterParamsLock.Unlock()

	// Create the HTTP/2 server that will be shared by both RPC and regular
	// duties. Doing it this way instead of listening via the server and gRPC
	// allows us to re-use the same port via ALPN. We can just tell the server
	// to serve a given conn and which handler to use.
	fws := &http2.Server{}

	// Shutdown coordination logic
	var shutdown uint32
	shutdownWg := &sync.WaitGroup{}

	for _, addr := range c.clusterListenerAddrs {
		shutdownWg.Add(1)

		// Force a local resolution to avoid data races
		laddr := addr

		// Start our listening loop
		go func() {
			defer shutdownWg.Done()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: starting listener", "listener_address", laddr)
			}

			// Create a TCP listener. We do this separately and specifically
			// with TCP so that we can set deadlines.
			tcpLn, err := net.ListenTCP("tcp", laddr)
			if err != nil {
				c.logger.Error("core/startClusterListener: error starting listener", "error", err)
				return
			}

			// Wrap the listener with TLS
			tlsLn := tls.NewListener(tcpLn, tlsConfig)
			defer tlsLn.Close()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: serving cluster requests", "cluster_listen_address", tlsLn.Addr())
			}

			for {
				if atomic.LoadUint32(&shutdown) > 0 {
					return
				}

				// Set the deadline for the accept call. If it passes we'll get
				// an error, causing us to check the condition at the top
				// again.
				tcpLn.SetDeadline(time.Now().Add(clusterListenerAcceptDeadline))

				// Accept the connection
				conn, err := tlsLn.Accept()
				if conn != nil {
					// Always defer although it may be closed ahead of time
					defer conn.Close()
				}
				if err != nil {
					continue
				}

				// Type assert to TLS connection and handshake to populate the
				// connection state
				tlsConn := conn.(*tls.Conn)
				err = tlsConn.Handshake()
				if err != nil {
					if c.logger.IsDebug() {
						c.logger.Debug("core: error handshaking cluster connection", "error", err)
					}
					if conn != nil {
						conn.Close()
					}
					continue
				}

				switch tlsConn.ConnectionState().NegotiatedProtocol {
				case "req_fw_sb-act_v1":
					if !ha {
						conn.Close()
						continue
					}

					c.logger.Trace("core: got req_fw_sb-act_v1 connection")
					go fws.ServeConn(conn, &http2.ServeConnOpts{
						Handler: c.rpcServer,
					})

				default:
					c.logger.Debug("core: unknown negotiated protocol on cluster port")
					conn.Close()
					continue
				}
			}
		}()
	}

	// This is in its own goroutine so that we don't block the main thread, and
	// thus we use atomic and channels to coordinate
	// However, because you can't query the status of a channel, we set a bool
	// here while we have the state lock to know whether to actually send a
	// shutdown (e.g. whether the channel will block). See issue #2083.
	c.clusterListenersRunning = true
	go func() {
		// If we get told to shut down...
		<-c.clusterListenerShutdownCh

		// Stop the RPC server
		c.logger.Info("core: shutting down forwarding rpc listeners")
		c.clusterParamsLock.Lock()
		c.rpcServer.Stop()
		c.rpcServer = nil
		c.clusterParamsLock.Unlock()
		c.logger.Info("core: forwarding rpc listeners stopped")

		// Set the shutdown flag. This will cause the listeners to shut down
		// within the deadline in clusterListenerAcceptDeadline
		atomic.StoreUint32(&shutdown, 1)

		// Wait for them all to shut down
		shutdownWg.Wait()
		c.logger.Info("core: rpc listeners successfully shut down")

		// Tell the main thread that shutdown is done.
		c.clusterListenerShutdownSuccessCh <- struct{}{}
	}()

	return nil
}

// refreshRequestForwardingConnection ensures that the client/transport are
// alive and that the current active address value matches the most
// recently-known address.
func (c *Core) refreshRequestForwardingConnection(clusterAddr string) error {
	c.logger.Trace("core: refreshing forwarding connection")
	defer c.logger.Trace("core: done refreshing forwarding connection")

	c.requestForwardingConnectionLock.Lock()
	defer c.requestForwardingConnectionLock.Unlock()

	// Clean things up first
	c.clearForwardingClients()

	// If we don't have anything to connect to, just return
	if clusterAddr == "" {
		return nil
	}

	clusterURL, err := url.Parse(clusterAddr)
	if err != nil {
		c.logger.Error("core: error parsing cluster address attempting to refresh forwarding connection", "error", err)
		return err
	}

	// Set up grpc forwarding handling
	// It's not really insecure, but we have to dial manually to get the
	// ALPN header right. It's just "insecure" because GRPC isn't managing
	// the TLS state.
	ctx, cancelFunc := context.WithCancel(context.Background())
	c.rpcClientConn, err = grpc.DialContext(ctx, clusterURL.Host,
		grpc.WithDialer(c.getGRPCDialer("req_fw_sb-act_v1", "", nil)),
		grpc.WithInsecure(), // it's not, we handle it in the dialer
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time: 2 * heartbeatInterval,
		}))
	if err != nil {
		cancelFunc()
		c.logger.Error("core: err setting up forwarding rpc client", "error", err)
		return err
	}
	c.rpcClientConnContext = ctx
	c.rpcClientConnCancelFunc = cancelFunc
	c.rpcForwardingClient = &forwardingClient{
		RequestForwardingClient: NewRequestForwardingClient(c.rpcClientConn),
		core:        c,
		echoTicker:  time.NewTicker(heartbeatInterval),
		echoContext: ctx,
	}
	c.rpcForwardingClient.startHeartbeat()

	return nil
}

func (c *Core) clearForwardingClients() {
	c.logger.Trace("core: clearing forwarding clients")
	defer c.logger.Trace("core: done clearing forwarding clients")

	if c.rpcClientConnCancelFunc != nil {
		c.rpcClientConnCancelFunc()
		c.rpcClientConnCancelFunc = nil
	}
	if c.rpcClientConn != nil {
		c.rpcClientConn.Close()
		c.rpcClientConn = nil
	}

	c.rpcClientConnContext = nil
	c.rpcForwardingClient = nil
}

// ForwardRequest forwards a given request to the active node and returns the
// response.
func (c *Core) ForwardRequest(req *http.Request) (int, http.Header, []byte, error) {
	c.requestForwardingConnectionLock.RLock()
	defer c.requestForwardingConnectionLock.RUnlock()

	if c.rpcForwardingClient == nil {
		return 0, nil, nil, ErrCannotForward
	}

	freq, err := forwarding.GenerateForwardedRequest(req)
	if err != nil {
		c.logger.Error("core: error creating forwarding RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error creating forwarding RPC request")
	}
	if freq == nil {
		c.logger.Error("core: got nil forwarding RPC request")
		return 0, nil, nil, fmt.Errorf("got nil forwarding RPC request")
	}
	resp, err := c.rpcForwardingClient.ForwardRequest(c.rpcClientConnContext, freq)
	if err != nil {
		c.logger.Error("core: error during forwarded RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error during forwarding RPC request")
	}

	var header http.Header
	if resp.HeaderEntries != nil {
		header = make(http.Header)
		for k, v := range resp.HeaderEntries {
			for _, j := range v.Values {
				header.Add(k, j)
			}
		}
	}

	return int(resp.StatusCode), header, resp.Body, nil
}

// getGRPCDialer is used to return a dialer that has the correct TLS
// configuration. Otherwise gRPC tries to be helpful and stomps all over our
// NextProtos.
func (c *Core) getGRPCDialer(alpnProto, serverName string, caCert *x509.Certificate) func(string, time.Duration) (net.Conn, error) {
	return func(addr string, timeout time.Duration) (net.Conn, error) {
		tlsConfig, err := c.ClusterTLSConfig()
		if err != nil {
			c.logger.Error("core: failed to get tls configuration", "error", err)
			return nil, err
		}
		if serverName != "" {
			tlsConfig.ServerName = serverName
		}
		if caCert != nil {
			pool := x509.NewCertPool()
			pool.AddCert(caCert)
			tlsConfig.RootCAs = pool
			tlsConfig.ClientCAs = pool
		}
		c.logger.Trace("core: creating rpc dialer", "host", tlsConfig.ServerName)

		tlsConfig.NextProtos = []string{alpnProto}
		dialer := &net.Dialer{
			Timeout: timeout,
		}
		return tls.DialWithDialer(dialer, "tcp", addr, tlsConfig)
	}
}

type forwardedRequestRPCServer struct {
	core    *Core
	handler http.Handler
}

func (s *forwardedRequestRPCServer) ForwardRequest(ctx context.Context, freq *forwarding.Request) (*forwarding.Response, error) {
	//s.core.logger.Trace("forwarding: serving rpc forwarded request")

	// Parse an http.Request out of it
	req, err := forwarding.ParseForwardedRequest(freq)
	if err != nil {
		return nil, err
	}

	// A very dummy response writer that doesn't follow normal semantics, just
	// lets you write a status code (last written wins) and a body. But it
	// meets the interface requirements.
	w := forwarding.NewRPCResponseWriter()

	s.handler.ServeHTTP(w, req)

	resp := &forwarding.Response{
		StatusCode: uint32(w.StatusCode()),
		Body:       w.Body().Bytes(),
	}

	header := w.Header()
	if header != nil {
		resp.HeaderEntries = make(map[string]*forwarding.HeaderEntry, len(header))
		for k, v := range header {
			resp.HeaderEntries[k] = &forwarding.HeaderEntry{
				Values: v,
			}
		}
	}

	return resp, nil
}

func (s *forwardedRequestRPCServer) Echo(ctx context.Context, in *EchoRequest) (*EchoReply, error) {
	if in.ClusterAddr != "" {
		s.core.clusterPeerClusterAddrsCache.Set(in.ClusterAddr, nil, 0)
	}
	return &EchoReply{
		Message: "pong",
	}, nil
}

type forwardingClient struct {
	RequestForwardingClient

	core *Core

	echoTicker  *time.Ticker
	echoContext context.Context
}

// NOTE: we also take advantage of gRPC's keepalive bits, but as we send data
// with these requests it's useful to keep this as well
func (c *forwardingClient) startHeartbeat() {
	go func() {
		tick := func() {
			c.core.stateLock.RLock()
			clusterAddr := c.core.clusterAddr
			c.core.stateLock.RUnlock()

			ctx, cancel := context.WithTimeout(c.echoContext, 2*time.Second)
			resp, err := c.RequestForwardingClient.Echo(ctx, &EchoRequest{
				Message:     "ping",
				ClusterAddr: clusterAddr,
			})
			cancel()
			if err != nil {
				c.core.logger.Debug("forwarding: error sending echo request to active node", "error", err)
				return
			}
			if resp == nil {
				c.core.logger.Debug("forwarding: empty echo response from active node")
				return
			}
			if resp.Message != "pong" {
				c.core.logger.Debug("forwarding: unexpected echo response from active node", "message", resp.Message)
				return
			}
			c.core.logger.Trace("forwarding: successful heartbeat")
		}

		tick()

		for {
			select {
			case <-c.echoContext.Done():
				c.echoTicker.Stop()
				return
			case <-c.echoTicker.C:
				tick()
			}
		}
	}()
}
<commit_msg>Log heartbeat stopping<commit_after>package vault

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/vault/helper/forwarding"
	"golang.org/x/net/context"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/keepalive"
)

const (
	clusterListenerAcceptDeadline = 500 * time.Millisecond
	heartbeatInterval             = 30 * time.Second
)

// Starts the listeners and servers necessary to handle forwarded requests
func (c *Core) startForwarding() error {
	c.logger.Trace("core: cluster listener setup function")
	defer c.logger.Trace("core: leaving cluster listener setup function")

	// Clean up in case we have transitioned from a client to a server
	c.requestForwardingConnectionLock.Lock()
	c.clearForwardingClients()
	c.requestForwardingConnectionLock.Unlock()

	// Resolve locally to avoid races
	ha := c.ha != nil

	// Get our TLS config
	tlsConfig, err := c.ClusterTLSConfig()
	if err != nil {
		c.logger.Error("core: failed to get tls configuration when starting forwarding", "error", err)
		return err
	}

	// The server supports all of the possible protos
	tlsConfig.NextProtos = []string{"h2", "req_fw_sb-act_v1"}

	// Create our RPC server and register the request handler server
	c.clusterParamsLock.Lock()

	if c.rpcServer != nil {
		c.logger.Warn("core: forwarding rpc server already running")
		return nil
	}

	c.rpcServer = grpc.NewServer(
		grpc.KeepaliveParams(keepalive.ServerParameters{
			Time: 2 * heartbeatInterval,
		}),
	)

	if ha && c.clusterHandler != nil {
		RegisterRequestForwardingServer(c.rpcServer, &forwardedRequestRPCServer{
			core:    c,
			handler: c.clusterHandler,
		})
	}
	c.clusterParamsLock.Unlock()

	// Create the HTTP/2 server that will be shared by both RPC and regular
	// duties. Doing it this way instead of listening via the server and gRPC
	// allows us to re-use the same port via ALPN. We can just tell the server
	// to serve a given conn and which handler to use.
	fws := &http2.Server{}

	// Shutdown coordination logic
	var shutdown uint32
	shutdownWg := &sync.WaitGroup{}

	for _, addr := range c.clusterListenerAddrs {
		shutdownWg.Add(1)

		// Force a local resolution to avoid data races
		laddr := addr

		// Start our listening loop
		go func() {
			defer shutdownWg.Done()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: starting listener", "listener_address", laddr)
			}

			// Create a TCP listener. We do this separately and specifically
			// with TCP so that we can set deadlines.
			tcpLn, err := net.ListenTCP("tcp", laddr)
			if err != nil {
				c.logger.Error("core/startClusterListener: error starting listener", "error", err)
				return
			}

			// Wrap the listener with TLS
			tlsLn := tls.NewListener(tcpLn, tlsConfig)
			defer tlsLn.Close()

			if c.logger.IsInfo() {
				c.logger.Info("core/startClusterListener: serving cluster requests", "cluster_listen_address", tlsLn.Addr())
			}

			for {
				if atomic.LoadUint32(&shutdown) > 0 {
					return
				}

				// Set the deadline for the accept call. If it passes we'll get
				// an error, causing us to check the condition at the top
				// again.
				tcpLn.SetDeadline(time.Now().Add(clusterListenerAcceptDeadline))

				// Accept the connection
				conn, err := tlsLn.Accept()
				if conn != nil {
					// Always defer although it may be closed ahead of time
					defer conn.Close()
				}
				if err != nil {
					continue
				}

				// Type assert to TLS connection and handshake to populate the
				// connection state
				tlsConn := conn.(*tls.Conn)
				err = tlsConn.Handshake()
				if err != nil {
					if c.logger.IsDebug() {
						c.logger.Debug("core: error handshaking cluster connection", "error", err)
					}
					if conn != nil {
						conn.Close()
					}
					continue
				}

				switch tlsConn.ConnectionState().NegotiatedProtocol {
				case "req_fw_sb-act_v1":
					if !ha {
						conn.Close()
						continue
					}

					c.logger.Trace("core: got req_fw_sb-act_v1 connection")
					go fws.ServeConn(conn, &http2.ServeConnOpts{
						Handler: c.rpcServer,
					})

				default:
					c.logger.Debug("core: unknown negotiated protocol on cluster port")
					conn.Close()
					continue
				}
			}
		}()
	}

	// This is in its own goroutine so that we don't block the main thread, and
	// thus we use atomic and channels to coordinate
	// However, because you can't query the status of a channel, we set a bool
	// here while we have the state lock to know whether to actually send a
	// shutdown (e.g. whether the channel will block). See issue #2083.
	c.clusterListenersRunning = true
	go func() {
		// If we get told to shut down...
		<-c.clusterListenerShutdownCh

		// Stop the RPC server
		c.logger.Info("core: shutting down forwarding rpc listeners")
		c.clusterParamsLock.Lock()
		c.rpcServer.Stop()
		c.rpcServer = nil
		c.clusterParamsLock.Unlock()
		c.logger.Info("core: forwarding rpc listeners stopped")

		// Set the shutdown flag. This will cause the listeners to shut down
		// within the deadline in clusterListenerAcceptDeadline
		atomic.StoreUint32(&shutdown, 1)

		// Wait for them all to shut down
		shutdownWg.Wait()
		c.logger.Info("core: rpc listeners successfully shut down")

		// Tell the main thread that shutdown is done.
		c.clusterListenerShutdownSuccessCh <- struct{}{}
	}()

	return nil
}

// refreshRequestForwardingConnection ensures that the client/transport are
// alive and that the current active address value matches the most
// recently-known address.
func (c *Core) refreshRequestForwardingConnection(clusterAddr string) error {
	c.logger.Trace("core: refreshing forwarding connection")
	defer c.logger.Trace("core: done refreshing forwarding connection")

	c.requestForwardingConnectionLock.Lock()
	defer c.requestForwardingConnectionLock.Unlock()

	// Clean things up first
	c.clearForwardingClients()

	// If we don't have anything to connect to, just return
	if clusterAddr == "" {
		return nil
	}

	clusterURL, err := url.Parse(clusterAddr)
	if err != nil {
		c.logger.Error("core: error parsing cluster address attempting to refresh forwarding connection", "error", err)
		return err
	}

	// Set up grpc forwarding handling
	// It's not really insecure, but we have to dial manually to get the
	// ALPN header right. It's just "insecure" because GRPC isn't managing
	// the TLS state.
	ctx, cancelFunc := context.WithCancel(context.Background())
	c.rpcClientConn, err = grpc.DialContext(ctx, clusterURL.Host,
		grpc.WithDialer(c.getGRPCDialer("req_fw_sb-act_v1", "", nil)),
		grpc.WithInsecure(), // it's not, we handle it in the dialer
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time: 2 * heartbeatInterval,
		}))
	if err != nil {
		cancelFunc()
		c.logger.Error("core: err setting up forwarding rpc client", "error", err)
		return err
	}
	c.rpcClientConnContext = ctx
	c.rpcClientConnCancelFunc = cancelFunc
	c.rpcForwardingClient = &forwardingClient{
		RequestForwardingClient: NewRequestForwardingClient(c.rpcClientConn),
		core:        c,
		echoTicker:  time.NewTicker(heartbeatInterval),
		echoContext: ctx,
	}
	c.rpcForwardingClient.startHeartbeat()

	return nil
}

func (c *Core) clearForwardingClients() {
	c.logger.Trace("core: clearing forwarding clients")
	defer c.logger.Trace("core: done clearing forwarding clients")

	if c.rpcClientConnCancelFunc != nil {
		c.rpcClientConnCancelFunc()
		c.rpcClientConnCancelFunc = nil
	}
	if c.rpcClientConn != nil {
		c.rpcClientConn.Close()
		c.rpcClientConn = nil
	}

	c.rpcClientConnContext = nil
	c.rpcForwardingClient = nil
}

// ForwardRequest forwards a given request to the active node and returns the
// response.
func (c *Core) ForwardRequest(req *http.Request) (int, http.Header, []byte, error) {
	c.requestForwardingConnectionLock.RLock()
	defer c.requestForwardingConnectionLock.RUnlock()

	if c.rpcForwardingClient == nil {
		return 0, nil, nil, ErrCannotForward
	}

	freq, err := forwarding.GenerateForwardedRequest(req)
	if err != nil {
		c.logger.Error("core: error creating forwarding RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error creating forwarding RPC request")
	}
	if freq == nil {
		c.logger.Error("core: got nil forwarding RPC request")
		return 0, nil, nil, fmt.Errorf("got nil forwarding RPC request")
	}
	resp, err := c.rpcForwardingClient.ForwardRequest(c.rpcClientConnContext, freq)
	if err != nil {
		c.logger.Error("core: error during forwarded RPC request", "error", err)
		return 0, nil, nil, fmt.Errorf("error during forwarding RPC request")
	}

	var header http.Header
	if resp.HeaderEntries != nil {
		header = make(http.Header)
		for k, v := range resp.HeaderEntries {
			for _, j := range v.Values {
				header.Add(k, j)
			}
		}
	}

	return int(resp.StatusCode), header, resp.Body, nil
}

// getGRPCDialer is used to return a dialer that has the correct TLS
// configuration. Otherwise gRPC tries to be helpful and stomps all over our
// NextProtos.
func (c *Core) getGRPCDialer(alpnProto, serverName string, caCert *x509.Certificate) func(string, time.Duration) (net.Conn, error) {
	return func(addr string, timeout time.Duration) (net.Conn, error) {
		tlsConfig, err := c.ClusterTLSConfig()
		if err != nil {
			c.logger.Error("core: failed to get tls configuration", "error", err)
			return nil, err
		}
		if serverName != "" {
			tlsConfig.ServerName = serverName
		}
		if caCert != nil {
			pool := x509.NewCertPool()
			pool.AddCert(caCert)
			tlsConfig.RootCAs = pool
			tlsConfig.ClientCAs = pool
		}
		c.logger.Trace("core: creating rpc dialer", "host", tlsConfig.ServerName)

		tlsConfig.NextProtos = []string{alpnProto}
		dialer := &net.Dialer{
			Timeout: timeout,
		}
		return tls.DialWithDialer(dialer, "tcp", addr, tlsConfig)
	}
}

type forwardedRequestRPCServer struct {
	core    *Core
	handler http.Handler
}

func (s *forwardedRequestRPCServer) ForwardRequest(ctx context.Context, freq *forwarding.Request) (*forwarding.Response, error) {
	//s.core.logger.Trace("forwarding: serving rpc forwarded request")

	// Parse an http.Request out of it
	req, err := forwarding.ParseForwardedRequest(freq)
	if err != nil {
		return nil, err
	}

	// A very dummy response writer that doesn't follow normal semantics, just
	// lets you write a status code (last written wins) and a body. But it
	// meets the interface requirements.
	w := forwarding.NewRPCResponseWriter()

	s.handler.ServeHTTP(w, req)

	resp := &forwarding.Response{
		StatusCode: uint32(w.StatusCode()),
		Body:       w.Body().Bytes(),
	}

	header := w.Header()
	if header != nil {
		resp.HeaderEntries = make(map[string]*forwarding.HeaderEntry, len(header))
		for k, v := range header {
			resp.HeaderEntries[k] = &forwarding.HeaderEntry{
				Values: v,
			}
		}
	}

	return resp, nil
}

func (s *forwardedRequestRPCServer) Echo(ctx context.Context, in *EchoRequest) (*EchoReply, error) {
	if in.ClusterAddr != "" {
		s.core.clusterPeerClusterAddrsCache.Set(in.ClusterAddr, nil, 0)
	}
	return &EchoReply{
		Message: "pong",
	}, nil
}

type forwardingClient struct {
	RequestForwardingClient

	core *Core

	echoTicker  *time.Ticker
	echoContext context.Context
}

// NOTE: we also take advantage of gRPC's keepalive bits, but as we send data
// with these requests it's useful to keep this as well
func (c *forwardingClient) startHeartbeat() {
	go func() {
		tick := func() {
			c.core.stateLock.RLock()
			clusterAddr := c.core.clusterAddr
			c.core.stateLock.RUnlock()

			ctx, cancel := context.WithTimeout(c.echoContext, 2*time.Second)
			resp, err := c.RequestForwardingClient.Echo(ctx, &EchoRequest{
				Message:     "ping",
				ClusterAddr: clusterAddr,
			})
			cancel()
			if err != nil {
				c.core.logger.Debug("forwarding: error sending echo request to active node", "error", err)
				return
			}
			if resp == nil {
				c.core.logger.Debug("forwarding: empty echo response from active node")
				return
			}
			if resp.Message != "pong" {
				c.core.logger.Debug("forwarding: unexpected echo response from active node", "message", resp.Message)
				return
			}
			c.core.logger.Trace("forwarding: successful heartbeat")
		}

		tick()

		for {
			select {
			case <-c.echoContext.Done():
				c.echoTicker.Stop()
				c.core.logger.Trace("forwarding: stopping heartbeating")
				return
			case <-c.echoTicker.C:
				tick()
			}
		}
	}()
}
<|endoftext|>


<commit_before>package vm

type Closure struct {
	function      *Thunk
	freeVariables []*Thunk
}

func (c Closure) Call(ts ...*Thunk) *Thunk {
	o := c.function.Eval()
	f, ok := o.(Callable)

	if !ok {
		return NotCallableError(o)
	}

	return f.Call(append(c.freeVariables, ts...)...)
}

func Partial(ts ...*Thunk) *Thunk {
	if len(ts) == 0 {
		return NumArgsError("partial", ">= 1")
	}

	return Normal(Closure{ts[0], ts[1:]})
}
<commit_msg>Fix out of range error<commit_after>package vm

type Closure struct {
	function      *Thunk
	freeVariables []*Thunk
}

func (c Closure) Call(ts ...*Thunk) *Thunk {
	o := c.function.Eval()
	f, ok := o.(Callable)

	if !ok {
		return NotCallableError(o)
	}

	return f.Call(append(c.freeVariables, ts...)...)
}

func Partial(ts ...*Thunk) *Thunk {
	switch len(ts) {
	case 0:
		return NumArgsError("partial", ">= 1")
	case 1:
		return ts[0]
	}

	return Normal(Closure{ts[0], ts[1:]})
}
<|endoftext|>


<commit_before>// Copyright 2017 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package vpython

import (
	"strings"

	"github.com/golang/protobuf/proto"
)

// Version is a version string. It must be updated any time the text protobufs
// advance in a non-backwards-compatible way.
//
// This version string is used in the generation of filenames, and must be
// filesystem-compatible.
const Version = "v1"

// Clone returns a deep clone of the supplied Environment.
//
// If e is nil, a non-nil empty Environment will be returned.
func (e *Environment) Clone() *Environment {
	if e == nil {
		return &Environment{}
	}
	return proto.Clone(e).(*Environment)
}

// IsZero returns true if this tag is a zero value.
func (t *Environment_Pep425Tag) IsZero() bool {
	return t == nil || (t.Version == "" && t.Abi == "" && t.Arch == "")
}

// TagString returns an underscore-separated string containing t's fields.
func (t *Environment_Pep425Tag) TagString() string {
	return strings.Join([]string{t.Version, t.Abi, t.Arch}, "-")
}

// HasABI returns true if t declares that it only works with a specific ABI.
func (t *Environment_Pep425Tag) HasABI() bool { return t.Abi != "none" }

// AnyArch returns true if t declares that it works on any architecture.
func (t *Environment_Pep425Tag) AnyArch() bool { return t.Arch != "any" }

// Count returns the number of populated fields in this tag.
func (t *Environment_Pep425Tag) Count() (v int) {
	if t.HasABI() {
		v++
	}
	if t.AnyArch() {
		v++
	}
	return
}

// Clone returns a deep clone of the supplied Spec.
//
// If e is nil, a non-nil empty Spec will be returned.
func (s *Spec) Clone() *Spec {
	if s == nil {
		return &Spec{}
	}
	return proto.Clone(s).(*Spec)
}
<commit_msg>vpython: Fix AnyArch.<commit_after>// Copyright 2017 The LUCI Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package vpython

import (
	"strings"

	"github.com/golang/protobuf/proto"
)

// Version is a version string. It must be updated any time the text protobufs
// advance in a non-backwards-compatible way.
//
// This version string is used in the generation of filenames, and must be
// filesystem-compatible.
const Version = "v1"

// Clone returns a deep clone of the supplied Environment.
//
// If e is nil, a non-nil empty Environment will be returned.
func (e *Environment) Clone() *Environment {
	if e == nil {
		return &Environment{}
	}
	return proto.Clone(e).(*Environment)
}

// IsZero returns true if this tag is a zero value.
func (t *Environment_Pep425Tag) IsZero() bool {
	return t == nil || (t.Version == "" && t.Abi == "" && t.Arch == "")
}

// TagString returns an underscore-separated string containing t's fields.
func (t *Environment_Pep425Tag) TagString() string {
	return strings.Join([]string{t.Version, t.Abi, t.Arch}, "-")
}

// HasABI returns true if t declares that it only works with a specific ABI.
func (t *Environment_Pep425Tag) HasABI() bool { return t.Abi != "none" }

// AnyArch returns true if t declares that it works on any architecture.
func (t *Environment_Pep425Tag) AnyArch() bool { return t.Arch == "any" }

// Count returns the number of populated fields in this tag.
func (t *Environment_Pep425Tag) Count() (v int) {
	if t.HasABI() {
		v++
	}
	if t.AnyArch() {
		v++
	}
	return
}

// Clone returns a deep clone of the supplied Spec.
//
// If e is nil, a non-nil empty Spec will be returned.
func (s *Spec) Clone() *Spec {
	if s == nil {
		return &Spec{}
	}
	return proto.Clone(s).(*Spec)
}
<|endoftext|>


<commit_before>package goble

/*
#include <xpc/xpc.h>
#include "xpc_wrapper.h"
*/
import "C"

import (
	"errors"
	"fmt"
	"log"
	r "reflect"
	"unsafe"
)

//
// minimal XPC support required for BLE
//

// a dictionary of things
type dict map[string]interface{}

func (d dict) GetString(k, defv string) string {
    if v := d[k]; v != nil {
        //log.Printf("GetString %s %#v\n", k, v)
        return v.(string)
    } else {
        //log.Printf("GetString %s default %#v\n", k, defv)
        return defv
    }
}

func (d dict) GetBytes(k string, defv []byte) []byte {
    if v := d[k]; v != nil {
        //log.Printf("GetBytes %s %#v\n", k, v)
        return v.([]byte)
    } else {
        //log.Printf("GetBytes %s default %#v\n", k, defv)
        return defv
    }
}

func (d dict) GetInt(k string, defv int64) int64 {
    if v := d[k]; v != nil {
        //log.Printf("GetString %s %#v\n", k, v)
        return v.(int64)
    } else {
        //log.Printf("GetString %s default %#v\n", k, defv)
        return defv
    }
}

func (d dict) GetUUID(k string) UUID {
    return GetUUID(d[k])
}

// an array of things
type array []interface{}

func (a array) GetUUID(k int) UUID {
    return GetUUID(a[k])
}

// a UUID
type UUID [16]byte

func (uuid UUID) String() string {
    return fmt.Sprintf("%x", [16]byte(uuid))
}

func GetUUID(v interface{}) UUID {
    if v == nil {
        return UUID{}
    }

    if uuid, ok := v.(UUID); ok {
        return uuid
    }

    if bytes, ok := v.([]byte); ok {
        uuid := UUID{}

        for i, b := range bytes {
            uuid[i] = b
        }

        return uuid
    }

    if bytes, ok := v.([]uint8); ok {
        uuid := UUID{}

        for i, b := range bytes {
            uuid[i] = b
        }

        return uuid
    }

    log.Fatalf("invalid type for UUID: %#v", v)
    return UUID{}
}

var (
	CONNECTION_INVALID     = errors.New("connection invalid")
	CONNECTION_INTERRUPTED = errors.New("connection interrupted")
	CONNECTION_TERMINATED  = errors.New("connection terminated")
)

type EventHandler func(event dict, err error)

// XpcConnect connects to the Blued service
// (with some trickery to pass around the event callback handler)
func XpcConnect(eh EventHandler) C.xpc_connection_t {
	return C.XpcConnectBlued((*(*unsafe.Pointer)(unsafe.Pointer(&eh))))
}

//export HandleXPCEvent
func HandleXPCEvent(event C.xpc_object_t, p unsafe.Pointer) {
	t := C.xpc_get_type(event)
	eh := (*(*EventHandler)(unsafe.Pointer(&p)))

	if t == C.TYPE_ERROR {
		if event == C.ERROR_CONNECTION_INVALID {
			// The client process on the other end of the connection has either
			// crashed or cancelled the connection. After receiving this error,
			// the connection is in an invalid state, and you do not need to
			// call xpc_connection_cancel(). Just tear down any associated state
			// here.
			//log.Println("connection invalid")
			eh(nil, CONNECTION_INVALID)
		} else if event == C.ERROR_CONNECTION_INTERRUPTED {
			//log.Println("connection interrupted")
			eh(nil, CONNECTION_INTERRUPTED)
		} else if event == C.ERROR_CONNECTION_TERMINATED {
			// Handle per-connection termination cleanup.
			//log.Println("connection terminated")
			eh(nil, CONNECTION_TERMINATED)
		} else {
			//log.Println("got some error", event)
			eh(nil, fmt.Errorf("%v", event))
		}
	} else {
		eh(xpcToGo(event).(dict), nil)
	}
}

// goToXpc converts a go object to an xpc object
func goToXpc(o interface{}) C.xpc_object_t {
	return valueToXpc(r.ValueOf(o))
}

// valueToXpc converts a go Value to an xpc object
//
// note that not all the types are supported, but only the subset required for Blued
func valueToXpc(val r.Value) C.xpc_object_t {
	if !val.IsValid() {
		return nil
	}

	var xv C.xpc_object_t

	switch val.Kind() {
	case r.Int, r.Int8, r.Int16, r.Int32, r.Int64:
		xv = C.xpc_int64_create(C.int64_t(val.Int()))

	case r.Uint, r.Uint8, r.Uint16, r.Uint32:
		xv = C.xpc_int64_create(C.int64_t(val.Uint()))

	case r.String:
		xv = C.xpc_string_create(C.CString(val.String()))

	case r.Map:
		xv = C.xpc_dictionary_create(nil, nil, 0)
		for _, k := range val.MapKeys() {
			v := valueToXpc(val.MapIndex(k))
			C.xpc_dictionary_set_value(xv, C.CString(k.String()), v)
			if v != nil {
				C.xpc_release(v)
			}
		}

	case r.Array, r.Slice:
		xv = C.xpc_array_create(nil, 0)
		l := val.Len()

		for i := 0; i < l; i++ {
			v := valueToXpc(val.Index(i))
			C.xpc_array_append_value(xv, v)
			if v != nil {
				C.xpc_release(v)
			}
		}

	case r.Interface, r.Ptr:
		xv = valueToXpc(val.Elem())

	default:
		log.Fatalf("unsupported %#v", val.String())
	}

	return xv
}

//export ArraySet
func ArraySet(u unsafe.Pointer, i C.int, v C.xpc_object_t) {
	a := *(*array)(u)
	a[i] = xpcToGo(v)
}

//export DictSet
func DictSet(u unsafe.Pointer, k *C.char, v C.xpc_object_t) {
	d := *(*dict)(u)
	d[C.GoString(k)] = xpcToGo(v)
}

// xpcToGo converts an xpc object to a go object
//
// note that not all the types are supported, but only the subset required for Blued
func xpcToGo(v C.xpc_object_t) interface{} {
	t := C.xpc_get_type(v)

	switch t {
	case C.TYPE_ARRAY:
		a := make(array, C.int(C.xpc_array_get_count(v)))
		C.XpcArrayApply(unsafe.Pointer(&a), v)
		return a

	case C.TYPE_DATA:
		return C.GoBytes(C.xpc_data_get_bytes_ptr(v), C.int(C.xpc_data_get_length(v)))

	case C.TYPE_DICT:
		d := make(dict)
		C.XpcDictApply(unsafe.Pointer(&d), v)
		return d

	case C.TYPE_INT64:
		return int64(C.xpc_int64_get_value(v))

	case C.TYPE_STRING:
		return C.GoString(C.xpc_string_get_string_ptr(v))

	case C.TYPE_UUID:
                a := [16]byte{}
		C.XpcUUIDGetBytes(unsafe.Pointer(&a), v)
		return UUID(a)

	default:
		log.Fatalf("unexpected type %#v, value %#v", t, v)
	}

	return nil
}
<commit_msg>added dict.Contains(key)<commit_after>package goble

/*
#include <xpc/xpc.h>
#include "xpc_wrapper.h"
*/
import "C"

import (
	"errors"
	"fmt"
	"log"
	r "reflect"
	"unsafe"
)

//
// minimal XPC support required for BLE
//

// a dictionary of things
type dict map[string]interface{}

func (d dict) Contains(k string) bool {
	_, ok := d[k]
	return ok
}

func (d dict) GetString(k, defv string) string {
	if v := d[k]; v != nil {
		//log.Printf("GetString %s %#v\n", k, v)
		return v.(string)
	} else {
		//log.Printf("GetString %s default %#v\n", k, defv)
		return defv
	}
}

func (d dict) GetBytes(k string, defv []byte) []byte {
	if v := d[k]; v != nil {
		//log.Printf("GetBytes %s %#v\n", k, v)
		return v.([]byte)
	} else {
		//log.Printf("GetBytes %s default %#v\n", k, defv)
		return defv
	}
}

func (d dict) GetInt(k string, defv int64) int64 {
	if v := d[k]; v != nil {
		//log.Printf("GetString %s %#v\n", k, v)
		return v.(int64)
	} else {
		//log.Printf("GetString %s default %#v\n", k, defv)
		return defv
	}
}

func (d dict) GetUUID(k string) UUID {
	return GetUUID(d[k])
}

// an array of things
type array []interface{}

func (a array) GetUUID(k int) UUID {
	return GetUUID(a[k])
}

// a UUID
type UUID [16]byte

func (uuid UUID) String() string {
	return fmt.Sprintf("%x", [16]byte(uuid))
}

func GetUUID(v interface{}) UUID {
	if v == nil {
		return UUID{}
	}

	if uuid, ok := v.(UUID); ok {
		return uuid
	}

	if bytes, ok := v.([]byte); ok {
		uuid := UUID{}

		for i, b := range bytes {
			uuid[i] = b
		}

		return uuid
	}

	if bytes, ok := v.([]uint8); ok {
		uuid := UUID{}

		for i, b := range bytes {
			uuid[i] = b
		}

		return uuid
	}

	log.Fatalf("invalid type for UUID: %#v", v)
	return UUID{}
}

var (
	CONNECTION_INVALID     = errors.New("connection invalid")
	CONNECTION_INTERRUPTED = errors.New("connection interrupted")
	CONNECTION_TERMINATED  = errors.New("connection terminated")
)

type EventHandler func(event dict, err error)

// XpcConnect connects to the Blued service
// (with some trickery to pass around the event callback handler)
func XpcConnect(eh EventHandler) C.xpc_connection_t {
	return C.XpcConnectBlued((*(*unsafe.Pointer)(unsafe.Pointer(&eh))))
}

//export HandleXPCEvent
func HandleXPCEvent(event C.xpc_object_t, p unsafe.Pointer) {
	t := C.xpc_get_type(event)
	eh := (*(*EventHandler)(unsafe.Pointer(&p)))

	if t == C.TYPE_ERROR {
		if event == C.ERROR_CONNECTION_INVALID {
			// The client process on the other end of the connection has either
			// crashed or cancelled the connection. After receiving this error,
			// the connection is in an invalid state, and you do not need to
			// call xpc_connection_cancel(). Just tear down any associated state
			// here.
			//log.Println("connection invalid")
			eh(nil, CONNECTION_INVALID)
		} else if event == C.ERROR_CONNECTION_INTERRUPTED {
			//log.Println("connection interrupted")
			eh(nil, CONNECTION_INTERRUPTED)
		} else if event == C.ERROR_CONNECTION_TERMINATED {
			// Handle per-connection termination cleanup.
			//log.Println("connection terminated")
			eh(nil, CONNECTION_TERMINATED)
		} else {
			//log.Println("got some error", event)
			eh(nil, fmt.Errorf("%v", event))
		}
	} else {
		eh(xpcToGo(event).(dict), nil)
	}
}

// goToXpc converts a go object to an xpc object
func goToXpc(o interface{}) C.xpc_object_t {
	return valueToXpc(r.ValueOf(o))
}

// valueToXpc converts a go Value to an xpc object
//
// note that not all the types are supported, but only the subset required for Blued
func valueToXpc(val r.Value) C.xpc_object_t {
	if !val.IsValid() {
		return nil
	}

	var xv C.xpc_object_t

	switch val.Kind() {
	case r.Int, r.Int8, r.Int16, r.Int32, r.Int64:
		xv = C.xpc_int64_create(C.int64_t(val.Int()))

	case r.Uint, r.Uint8, r.Uint16, r.Uint32:
		xv = C.xpc_int64_create(C.int64_t(val.Uint()))

	case r.String:
		xv = C.xpc_string_create(C.CString(val.String()))

	case r.Map:
		xv = C.xpc_dictionary_create(nil, nil, 0)
		for _, k := range val.MapKeys() {
			v := valueToXpc(val.MapIndex(k))
			C.xpc_dictionary_set_value(xv, C.CString(k.String()), v)
			if v != nil {
				C.xpc_release(v)
			}
		}

	case r.Array, r.Slice:
		xv = C.xpc_array_create(nil, 0)
		l := val.Len()

		for i := 0; i < l; i++ {
			v := valueToXpc(val.Index(i))
			C.xpc_array_append_value(xv, v)
			if v != nil {
				C.xpc_release(v)
			}
		}

	case r.Interface, r.Ptr:
		xv = valueToXpc(val.Elem())

	default:
		log.Fatalf("unsupported %#v", val.String())
	}

	return xv
}

//export ArraySet
func ArraySet(u unsafe.Pointer, i C.int, v C.xpc_object_t) {
	a := *(*array)(u)
	a[i] = xpcToGo(v)
}

//export DictSet
func DictSet(u unsafe.Pointer, k *C.char, v C.xpc_object_t) {
	d := *(*dict)(u)
	d[C.GoString(k)] = xpcToGo(v)
}

// xpcToGo converts an xpc object to a go object
//
// note that not all the types are supported, but only the subset required for Blued
func xpcToGo(v C.xpc_object_t) interface{} {
	t := C.xpc_get_type(v)

	switch t {
	case C.TYPE_ARRAY:
		a := make(array, C.int(C.xpc_array_get_count(v)))
		C.XpcArrayApply(unsafe.Pointer(&a), v)
		return a

	case C.TYPE_DATA:
		return C.GoBytes(C.xpc_data_get_bytes_ptr(v), C.int(C.xpc_data_get_length(v)))

	case C.TYPE_DICT:
		d := make(dict)
		C.XpcDictApply(unsafe.Pointer(&d), v)
		return d

	case C.TYPE_INT64:
		return int64(C.xpc_int64_get_value(v))

	case C.TYPE_STRING:
		return C.GoString(C.xpc_string_get_string_ptr(v))

	case C.TYPE_UUID:
		a := [16]byte{}
		C.XpcUUIDGetBytes(unsafe.Pointer(&a), v)
		return UUID(a)

	default:
		log.Fatalf("unexpected type %#v, value %#v", t, v)
	}

	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 y

import (
	"bytes"
	"encoding/binary"
	"hash/crc32"
	"math"
	"os"
	"sync"

	"github.com/pkg/errors"
)

// ErrEOF indicates an end of file when trying to read from a memory mapped file
// and encountering the end of slice.
var ErrEOF = errors.New("End of mapped region")

const (
	// Sync indicates that O_DSYNC should be set on the underlying file,
	// ensuring that data writes do not return until the data is flushed
	// to disk.
	Sync = 1 << iota
	// ReadOnly opens the underlying file on a read-only basis.
	ReadOnly
)

var (
	// This is O_DSYNC (datasync) on platforms that support it -- see file_unix.go
	datasyncFileFlag = 0x0

	// CastagnoliCrcTable is a CRC32 polynomial table
	CastagnoliCrcTable = crc32.MakeTable(crc32.Castagnoli)
)

// OpenExistingFile opens an existing file, errors if it doesn't exist.
func OpenExistingFile(filename string, flags uint32) (*os.File, error) {
	openFlags := os.O_RDWR
	if flags&ReadOnly != 0 {
		openFlags = os.O_RDONLY
	}

	if flags&Sync != 0 {
		openFlags |= datasyncFileFlag
	}
	return os.OpenFile(filename, openFlags, 0)
}

// CreateSyncedFile creates a new file (using O_EXCL), errors if it already existed.
func CreateSyncedFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE | os.O_EXCL
	if sync {
		flags |= datasyncFileFlag
	}
	return os.OpenFile(filename, flags, 0666)
}

// OpenSyncedFile creates the file if one doesn't exist.
func OpenSyncedFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE
	if sync {
		flags |= datasyncFileFlag
	}
	return os.OpenFile(filename, flags, 0666)
}

// OpenTruncFile opens the file with O_RDWR | O_CREATE | O_TRUNC
func OpenTruncFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE | os.O_TRUNC
	if sync {
		flags |= datasyncFileFlag
	}
	return os.OpenFile(filename, flags, 0666)
}

// SafeCopy does append(a[:0], src...).
func SafeCopy(a []byte, src []byte) []byte {
	return append(a[:0], src...)
}

// Copy copies a byte slice and returns the copied slice.
func Copy(a []byte) []byte {
	b := make([]byte, len(a))
	copy(b, a)
	return b
}

// KeyWithTs generates a new key by appending ts to key.
func KeyWithTs(key []byte, ts uint64) []byte {
	out := make([]byte, len(key)+8)
	copy(out, key)
	binary.BigEndian.PutUint64(out[len(key):], math.MaxUint64-ts)
	return out
}

// ParseTs parses the timestamp from the key bytes.
func ParseTs(key []byte) uint64 {
	if len(key) <= 8 {
		return 0
	}
	return math.MaxUint64 - binary.BigEndian.Uint64(key[len(key)-8:])
}

// CompareKeys checks the key without timestamp and checks the timestamp if keyNoTs
// is same.
// a<timestamp> would be sorted higher than aa<timestamp> if we use bytes.compare
// All keys should have timestamp.
func CompareKeys(key1 []byte, key2 []byte) int {
	AssertTrue(len(key1) > 8 && len(key2) > 8)
	if cmp := bytes.Compare(key1[:len(key1)-8], key2[:len(key2)-8]); cmp != 0 {
		return cmp
	}
	return bytes.Compare(key1[len(key1)-8:], key2[len(key2)-8:])
}

// ParseKey parses the actual key from the key bytes.
func ParseKey(key []byte) []byte {
	if key == nil {
		return nil
	}

	AssertTruef(len(key) > 8, "key=%q", key)
	return key[:len(key)-8]
}

// SameKey checks for key equality ignoring the version timestamp suffix.
func SameKey(src, dst []byte) bool {
	if len(src) != len(dst) {
		return false
	}
	return bytes.Equal(ParseKey(src), ParseKey(dst))
}

// Slice holds a reusable buf, will reallocate if you request a larger size than ever before.
// One problem is with n distinct sizes in random order it'll reallocate log(n) times.
type Slice struct {
	buf []byte
}

// Resize reuses the Slice's buffer (or makes a new one) and returns a slice in that buffer of
// length sz.
func (s *Slice) Resize(sz int) []byte {
	if cap(s.buf) < sz {
		s.buf = make([]byte, sz)
	}
	return s.buf[0:sz]
}

// Closer holds the two things we need to close a goroutine and wait for it to finish: a chan
// to tell the goroutine to shut down, and a WaitGroup with which to wait for it to finish shutting
// down.
type Closer struct {
	closed  chan struct{}
	waiting sync.WaitGroup
}

// NewCloser constructs a new Closer, with an initial count on the WaitGroup.
func NewCloser(initial int) *Closer {
	ret := &Closer{closed: make(chan struct{})}
	ret.waiting.Add(initial)
	return ret
}

// AddRunning Add()'s delta to the WaitGroup.
func (lc *Closer) AddRunning(delta int) {
	lc.waiting.Add(delta)
}

// Signal signals the HasBeenClosed signal.
func (lc *Closer) Signal() {
	close(lc.closed)
}

// HasBeenClosed gets signaled when Signal() is called.
func (lc *Closer) HasBeenClosed() <-chan struct{} {
	return lc.closed
}

// Done calls Done() on the WaitGroup.
func (lc *Closer) Done() {
	lc.waiting.Done()
}

// Wait waits on the WaitGroup.  (It waits for NewCloser's initial value, AddRunning, and Done
// calls to balance out.)
func (lc *Closer) Wait() {
	lc.waiting.Wait()
}

// SignalAndWait calls Signal(), then Wait().
func (lc *Closer) SignalAndWait() {
	lc.Signal()
	lc.Wait()
}
<commit_msg>reduce allocations by avoiding AssertTruef (#597)<commit_after>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 y

import (
	"bytes"
	"encoding/binary"
	"hash/crc32"
	"math"
	"os"
	"sync"

	"github.com/pkg/errors"
)

// ErrEOF indicates an end of file when trying to read from a memory mapped file
// and encountering the end of slice.
var ErrEOF = errors.New("End of mapped region")

const (
	// Sync indicates that O_DSYNC should be set on the underlying file,
	// ensuring that data writes do not return until the data is flushed
	// to disk.
	Sync = 1 << iota
	// ReadOnly opens the underlying file on a read-only basis.
	ReadOnly
)

var (
	// This is O_DSYNC (datasync) on platforms that support it -- see file_unix.go
	datasyncFileFlag = 0x0

	// CastagnoliCrcTable is a CRC32 polynomial table
	CastagnoliCrcTable = crc32.MakeTable(crc32.Castagnoli)
)

// OpenExistingFile opens an existing file, errors if it doesn't exist.
func OpenExistingFile(filename string, flags uint32) (*os.File, error) {
	openFlags := os.O_RDWR
	if flags&ReadOnly != 0 {
		openFlags = os.O_RDONLY
	}

	if flags&Sync != 0 {
		openFlags |= datasyncFileFlag
	}
	return os.OpenFile(filename, openFlags, 0)
}

// CreateSyncedFile creates a new file (using O_EXCL), errors if it already existed.
func CreateSyncedFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE | os.O_EXCL
	if sync {
		flags |= datasyncFileFlag
	}
	return os.OpenFile(filename, flags, 0666)
}

// OpenSyncedFile creates the file if one doesn't exist.
func OpenSyncedFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE
	if sync {
		flags |= datasyncFileFlag
	}
	return os.OpenFile(filename, flags, 0666)
}

// OpenTruncFile opens the file with O_RDWR | O_CREATE | O_TRUNC
func OpenTruncFile(filename string, sync bool) (*os.File, error) {
	flags := os.O_RDWR | os.O_CREATE | os.O_TRUNC
	if sync {
		flags |= datasyncFileFlag
	}
	return os.OpenFile(filename, flags, 0666)
}

// SafeCopy does append(a[:0], src...).
func SafeCopy(a []byte, src []byte) []byte {
	return append(a[:0], src...)
}

// Copy copies a byte slice and returns the copied slice.
func Copy(a []byte) []byte {
	b := make([]byte, len(a))
	copy(b, a)
	return b
}

// KeyWithTs generates a new key by appending ts to key.
func KeyWithTs(key []byte, ts uint64) []byte {
	out := make([]byte, len(key)+8)
	copy(out, key)
	binary.BigEndian.PutUint64(out[len(key):], math.MaxUint64-ts)
	return out
}

// ParseTs parses the timestamp from the key bytes.
func ParseTs(key []byte) uint64 {
	if len(key) <= 8 {
		return 0
	}
	return math.MaxUint64 - binary.BigEndian.Uint64(key[len(key)-8:])
}

// CompareKeys checks the key without timestamp and checks the timestamp if keyNoTs
// is same.
// a<timestamp> would be sorted higher than aa<timestamp> if we use bytes.compare
// All keys should have timestamp.
func CompareKeys(key1 []byte, key2 []byte) int {
	AssertTrue(len(key1) > 8 && len(key2) > 8)
	if cmp := bytes.Compare(key1[:len(key1)-8], key2[:len(key2)-8]); cmp != 0 {
		return cmp
	}
	return bytes.Compare(key1[len(key1)-8:], key2[len(key2)-8:])
}

// ParseKey parses the actual key from the key bytes.
func ParseKey(key []byte) []byte {
	if key == nil {
		return nil
	}

	AssertTrue(len(key) > 8)
	return key[:len(key)-8]
}

// SameKey checks for key equality ignoring the version timestamp suffix.
func SameKey(src, dst []byte) bool {
	if len(src) != len(dst) {
		return false
	}
	return bytes.Equal(ParseKey(src), ParseKey(dst))
}

// Slice holds a reusable buf, will reallocate if you request a larger size than ever before.
// One problem is with n distinct sizes in random order it'll reallocate log(n) times.
type Slice struct {
	buf []byte
}

// Resize reuses the Slice's buffer (or makes a new one) and returns a slice in that buffer of
// length sz.
func (s *Slice) Resize(sz int) []byte {
	if cap(s.buf) < sz {
		s.buf = make([]byte, sz)
	}
	return s.buf[0:sz]
}

// Closer holds the two things we need to close a goroutine and wait for it to finish: a chan
// to tell the goroutine to shut down, and a WaitGroup with which to wait for it to finish shutting
// down.
type Closer struct {
	closed  chan struct{}
	waiting sync.WaitGroup
}

// NewCloser constructs a new Closer, with an initial count on the WaitGroup.
func NewCloser(initial int) *Closer {
	ret := &Closer{closed: make(chan struct{})}
	ret.waiting.Add(initial)
	return ret
}

// AddRunning Add()'s delta to the WaitGroup.
func (lc *Closer) AddRunning(delta int) {
	lc.waiting.Add(delta)
}

// Signal signals the HasBeenClosed signal.
func (lc *Closer) Signal() {
	close(lc.closed)
}

// HasBeenClosed gets signaled when Signal() is called.
func (lc *Closer) HasBeenClosed() <-chan struct{} {
	return lc.closed
}

// Done calls Done() on the WaitGroup.
func (lc *Closer) Done() {
	lc.waiting.Done()
}

// Wait waits on the WaitGroup.  (It waits for NewCloser's initial value, AddRunning, and Done
// calls to balance out.)
func (lc *Closer) Wait() {
	lc.waiting.Wait()
}

// SignalAndWait calls Signal(), then Wait().
func (lc *Closer) SignalAndWait() {
	lc.Signal()
	lc.Wait()
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package watcher

import (
	"errors"
	"expvar"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"syscall"
	"testing"
	"time"

	"github.com/kylelemons/godebug/pretty"
)

func TestLogWatcher(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping log watcher test in short mode")
	}

	workdir, err := ioutil.TempDir("", "log_watcher_test")
	if err != nil {
		t.Fatalf("could not create temporary working directory: %s", err)
	}

	defer func() {
		err := os.RemoveAll(workdir)
		if err != nil {
			t.Fatalf("could not remove temp dir %s: %s:", workdir, err)
		}
	}()

	w, err := NewLogWatcher()
	if err != nil {
		t.Fatalf("couldn't create a watcher: %s\n", err)
	}
	defer w.Close()

	w.Add(workdir)
	f, err := os.Create(filepath.Join(workdir, "logfile"))
	if err != nil {
		t.Fatalf("couldn't make a logfile in temp dir: %s\n", err)
	}
	select {
	case e := <-w.Events():
		switch e := e.(type) {
		case CreateEvent:
			if e.Pathname != filepath.Join(workdir, "logfile") {
				t.Errorf("create doesn't match")
			}
		default:
			t.Errorf("Wrong event type: %q", e)
		}
	case <-time.After(10 * time.Millisecond):
		t.Errorf("didn't receive create message")
	}
	f.WriteString("hi")
	f.Close()
	select {
	case e := <-w.Events():
		switch e := e.(type) {
		case UpdateEvent:
			if e.Pathname != filepath.Join(workdir, "logfile") {
				t.Errorf("update doesn't match")
			}
		default:
			t.Errorf("Wrong event type: %q", e)
		}
	case <-time.After(10 * time.Millisecond):
		t.Errorf("didn't receive update message")
	}
	os.Chmod(filepath.Join(workdir, "logfile"), os.ModePerm)
	select {
	case e := <-w.Events():
		t.Errorf("no event expected, got %q", e)
	case <-time.After(10 * time.Millisecond):
	}
	os.Remove(filepath.Join(workdir, "logfile"))
	select {
	case e := <-w.Events():
		switch e := e.(type) {
		case DeleteEvent:
			if e.Pathname != filepath.Join(workdir, "logfile") {
				t.Errorf("delete doesn't match")
			}
		default:
			t.Errorf("Wrong event type: %q", e)
		}
	case <-time.After(10 * time.Millisecond):
		t.Errorf("didn't receive delete message")
	}
}

// This test may be OS specific; possibly break it out to a file with build tags.
func TestNewLogWatcherError(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping log watcher test in short mode")
	}

	var rLimit syscall.Rlimit
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil {
		t.Fatalf("coulnd't get rlimit: %s", err)
	}
	var zero = rLimit
	zero.Cur = 0
	if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &zero); err != nil {
		t.Fatalf("couldn't set rlimit: %s", err)
	}
	_, err := NewLogWatcher()
	if err == nil {
		t.Errorf("didn't fail as expected")
	}
	//t.Logf("expected error: %s", err)
	if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil {
		t.Fatalf("couldn't reset rlimit: %s", err)
	}
}

func TestLogWatcherAddError(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping log watcher test in short mode")
	}

	workdir, err := ioutil.TempDir("", "log_watcher_test")
	if err != nil {
		t.Fatalf("could not create temporary working directory: %s", err)
	}

	defer func() {
		err := os.RemoveAll(workdir)
		if err != nil {
			t.Fatalf("could not remove temp dir %s: %s:", workdir, err)
		}
	}()

	w, err := NewLogWatcher()
	if err != nil {
		t.Fatalf("couldn't create a watcher: %s\n", err)
	}
	defer w.Close()

	filename := filepath.Join(workdir, "test")
	if _, err := os.Create(filename); err != nil {
		t.Fatalf("couldn't create file: %s", err)
	}
	if err := os.Chmod(filename, 0); err != nil {
		t.Fatalf("couldn't chmod file: %s", err)
	}
	err = w.Add(filename)
	if err == nil {
		t.Errorf("didn't fail to add file")
	}
	//t.Logf("error: %s", err)
	if err := os.Chmod(filename, 0777); err != nil {
		t.Fatalf("coulnd't reset file perms: %s", err)
	}
}

func TestWatcherErrors(t *testing.T) {
	orig, err := strconv.ParseInt(expvar.Get("log_watcher_error_count").String(), 10, 64)
	if err != nil {
		t.Fatalf("couldn't convert expvar %q", expvar.Get("log_watcher_error_count").String())
	}
	w, err := NewLogWatcher()
	if err != nil {
		t.Fatalf("couldn't create a watcher")
	}
	w.Errors <- errors.New("test error")
	w.Close()
	diff := pretty.Compare(strconv.FormatInt(orig+1, 10), expvar.Get("log_watcher_error_count").String())
	if len(diff) > 0 {
		t.Errorf("log watcher error count doens't match:\n%s", diff)
	}
}
<commit_msg>Improve log message and check for error in log watcher test.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package watcher

import (
	"errors"
	"expvar"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"syscall"
	"testing"
	"time"

	"github.com/kylelemons/godebug/pretty"
)

func TestLogWatcher(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping log watcher test in short mode")
	}

	workdir, err := ioutil.TempDir("", "log_watcher_test")
	if err != nil {
		t.Fatalf("could not create temporary working directory: %s", err)
	}

	defer func() {
		err := os.RemoveAll(workdir)
		if err != nil {
			t.Fatalf("could not remove temp dir %s: %s:", workdir, err)
		}
	}()

	w, err := NewLogWatcher()
	if err != nil {
		t.Fatalf("couldn't create a watcher: %s\n", err)
	}
	defer w.Close()

	w.Add(workdir)
	f, err := os.Create(filepath.Join(workdir, "logfile"))
	if err != nil {
		t.Fatalf("couldn't make a logfile in temp dir: %s\n", err)
	}
	select {
	case e := <-w.Events():
		switch e := e.(type) {
		case CreateEvent:
			if e.Pathname != filepath.Join(workdir, "logfile") {
				t.Errorf("create doesn't match")
			}
		default:
			t.Errorf("Wrong event type: %q", e)
		}
	case <-time.After(10 * time.Millisecond):
		t.Errorf("didn't receive create message")
	}
	f.WriteString("hi")
	f.Close()
	select {
	case e := <-w.Events():
		switch e := e.(type) {
		case UpdateEvent:
			if e.Pathname != filepath.Join(workdir, "logfile") {
				t.Errorf("update doesn't match")
			}
		default:
			t.Errorf("Wrong event type: %q", e)
		}
	case <-time.After(10 * time.Millisecond):
		t.Errorf("didn't receive update message")
	}
	os.Chmod(filepath.Join(workdir, "logfile"), os.ModePerm)
	select {
	case e := <-w.Events():
		t.Errorf("no event expected, got %q", e)
	case <-time.After(10 * time.Millisecond):
	}
	os.Remove(filepath.Join(workdir, "logfile"))
	select {
	case e := <-w.Events():
		switch e := e.(type) {
		case DeleteEvent:
			if e.Pathname != filepath.Join(workdir, "logfile") {
				t.Errorf("delete doesn't match")
			}
		default:
			t.Errorf("Wrong event type: %q", e)
		}
	case <-time.After(10 * time.Millisecond):
		t.Errorf("didn't receive delete message")
	}
}

// This test may be OS specific; possibly break it out to a file with build tags.
func TestNewLogWatcherError(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping log watcher test in short mode")
	}

	var rLimit syscall.Rlimit
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil {
		t.Fatalf("coulnd't get rlimit: %s", err)
	}
	var zero = rLimit
	zero.Cur = 0
	if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &zero); err != nil {
		t.Fatalf("couldn't set rlimit: %s", err)
	}
	_, err := NewLogWatcher()
	if err == nil {
		t.Errorf("didn't fail as expected")
	}
	//t.Logf("expected error: %s", err)
	if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil {
		t.Fatalf("couldn't reset rlimit: %s", err)
	}
}

func TestLogWatcherAddError(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping log watcher test in short mode")
	}

	workdir, err := ioutil.TempDir("", "log_watcher_test")
	if err != nil {
		t.Fatalf("could not create temporary working directory: %s", err)
	}

	defer func() {
		err := os.RemoveAll(workdir)
		if err != nil {
			t.Fatalf("could not remove temp dir %s: %s:", workdir, err)
		}
	}()

	w, err := NewLogWatcher()
	if err != nil {
		t.Fatalf("couldn't create a watcher: %s\n", err)
	}
	defer w.Close()

	filename := filepath.Join(workdir, "test")
	if _, err := os.Create(filename); err != nil {
		t.Fatalf("couldn't create file: %s", err)
	}
	if err := os.Chmod(filename, 0); err != nil {
		t.Fatalf("couldn't chmod file: %s", err)
	}
	err = w.Add(filename)
	if err == nil {
		t.Errorf("didn't fail to add file")
	}
	//t.Logf("error: %s", err)
	if err := os.Chmod(filename, 0777); err != nil {
		t.Fatalf("coulnd't reset file perms: %s", err)
	}
}

func TestWatcherErrors(t *testing.T) {
	orig, err := strconv.ParseInt(expvar.Get("log_watcher_error_count").String(), 10, 64)
	if err != nil {
		t.Fatalf("couldn't convert expvar %q", expvar.Get("log_watcher_error_count").String())
	}
	w, err := NewLogWatcher()
	if err != nil {
		t.Fatalf("couldn't create a watcher")
	}
	w.Errors <- errors.New("just a test, not really an error")
	if err := w.Close(); err != nil {
		t.Fatalf("watcher close failed: %q", err)
	}
	diff := pretty.Compare(strconv.FormatInt(orig+1, 10), expvar.Get("log_watcher_error_count").String())
	if len(diff) > 0 {
		t.Errorf("log watcher error count doens't match:\n%s", diff)
	}
}
<|endoftext|>


<commit_before>// Package zlmgo provides Go bindings to the Zen License Manager (ZLM).
package zlmgo

/*
#cgo CFLAGS: -I/usr/local/include
#cgo LDFLAGS: /usr/local/lib/libzlm.a

#include <stdlib.h>
#include <zlm.h>

char zlmgo_errbuf_array[ZLM_ERRBUF];
char* zlmgo_errbuf = zlmgo_errbuf_array; // hack around cgo warning
*/
import "C"

import (
	"errors"
	"runtime"
	"unsafe"
)

// License wraps a ZLM license object.
type License struct {
	l *C.ZlmLicense
}

// LicenseNew returns a new license object (panics if not enough memory is available).
func LicenseNew() *License {
	license := &License{C.zlm_license_new(C.zlmgo_errbuf)}
	if license.l == nil {
		panic(C.GoString(C.zlmgo_errbuf))
	}
	runtime.SetFinalizer(license, (*License).free)
	return license
}

// Get wraps zlm_license_get(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Get(product, version, argv0, path, licenseString string) error {
	var cProduct, cVersion, cArgv0, cPath, cLicenseString *C.char
	// convert argument to C strings
	if product != "" {
		cProduct = C.CString(product)
	}
	if version != "" {
		cVersion = C.CString(version)
	}
	if argv0 != "" {
		cArgv0 = C.CString(argv0)
	}
	if path != "" {
		cPath = C.CString(path)
	}
	if licenseString != "" {
		cLicenseString = C.CString(licenseString)
	}
	// call actual method
	if C.zlm_license_get(license.l, cProduct, cVersion, cArgv0, cPath, cLicenseString, C.zlmgo_errbuf) != C.ZLM_OK {
		return errors.New(C.GoString(C.zlmgo_errbuf))
	}
	return nil
}

func (license *License) free() {
	C.zlm_license_free(license.l)
}

// Product wraps zlm_license_product(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Product() string {
	return C.GoString(C.zlm_license_product(license.l))
}

// Expiry wraps zlm_license_expiry(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Expiry() string {
	return C.GoString(C.zlm_license_expiry(license.l))
}

// ExpiryDays wraps zlm_license_expiry_days(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) ExpiryDays() int {
	return int(C.zlm_license_expiry_days(license.l))
}

// Customer wraps zlm_license_customer(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Customer() string {
	return C.GoString(C.zlm_license_customer(license.l))
}

// Userdata wraps zlm_license_userdata(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Userdata() string {
	return C.GoString(C.zlm_license_userdata(license.l))
}

// UserdataUnescaped wraps zlm_license_userdata_unescaped(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) UserdataUnescaped() string {
	return C.GoString(C.zlm_license_userdata_unescaped(license.l))
}

// Next wraps zlm_license_next(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Next() error {
	if C.zlm_license_next(license.l, C.zlmgo_errbuf) != C.ZLM_OK {
		return errors.New(C.GoString(C.zlmgo_errbuf))
	}
	return nil
}

// Version wraps zlm_version(), see
// https://zenlicensemanager.com/documentation/#API
func Version() string {
	return C.GoString(C.zlm_version())
}

// HostidJSON wraps zlm_hostid_json(), see
// https://zenlicensemanager.com/documentation/#API
func HostidJSON() (string, error) {
	cs := C.zlm_hostid_json(C.zlmgo_errbuf)
	if cs == nil {
		return "", errors.New(C.GoString(C.zlmgo_errbuf))
	}
	hostid := C.GoString(cs)
	C.zlm_free(unsafe.Pointer(cs))
	return hostid, nil
}

// CheckA wraps zlm_license_check_a(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckA() {
	C.zlm_license_check_a(license.l)
}

// CheckB wraps zlm_license_check_b(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckB() {
	C.zlm_license_check_b(license.l)
}

// CheckC wraps zlm_license_check_c(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckC() {
	C.zlm_license_check_c(license.l)
}

// CheckD wraps zlm_license_check_d(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckD() {
	C.zlm_license_check_d(license.l)
}

// CheckE wraps zlm_license_check_e(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckE() {
	C.zlm_license_check_e(license.l)
}

// CheckF wraps zlm_license_check_f(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckF() {
	C.zlm_license_check_f(license.l)
}

// CheckG wraps zlm_license_check_g(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckG() {
	C.zlm_license_check_g(license.l)
}

// CheckH wraps zlm_license_check_h(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckH() {
	C.zlm_license_check_h(license.l)
}

// CheckI wraps zlm_license_check_i(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckI() {
	C.zlm_license_check_i(license.l)
}

// CheckJ wraps zlm_license_check_j(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckJ() {
	C.zlm_license_check_j(license.l)
}

// CheckK wraps zlm_license_check_k(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckK() {
	C.zlm_license_check_k(license.l)
}

// CheckL wraps zlm_license_check_l(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckL() {
	C.zlm_license_check_l(license.l)
}

// CheckM wraps zlm_license_check_m(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckM() {
	C.zlm_license_check_m(license.l)
}

// CheckN wraps zlm_license_check_n(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckN() {
	C.zlm_license_check_n(license.l)
}

// CheckO wraps zlm_license_check_o(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckO() {
	C.zlm_license_check_o(license.l)
}

// CheckP wraps zlm_license_check_p(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckP() {
	C.zlm_license_check_p(license.l)
}

// CheckQ wraps zlm_license_check_q(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckQ() {
	C.zlm_license_check_q(license.l)
}

// CheckR wraps zlm_license_check_r(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckR() {
	C.zlm_license_check_r(license.l)
}
<commit_msg>add license.NotExpired()<commit_after>// Package zlmgo provides Go bindings to the Zen License Manager (ZLM).
package zlmgo

/*
#cgo CFLAGS: -I/usr/local/include
#cgo LDFLAGS: /usr/local/lib/libzlm.a

#include <stdlib.h>
#include <zlm.h>

char zlmgo_errbuf_array[ZLM_ERRBUF];
char* zlmgo_errbuf = zlmgo_errbuf_array; // hack around cgo warning
*/
import "C"

import (
	"errors"
	"runtime"
	"unsafe"
)

// License wraps a ZLM license object.
type License struct {
	l *C.ZlmLicense
}

// LicenseNew returns a new license object (panics if not enough memory is available).
func LicenseNew() *License {
	license := &License{C.zlm_license_new(C.zlmgo_errbuf)}
	if license.l == nil {
		panic(C.GoString(C.zlmgo_errbuf))
	}
	runtime.SetFinalizer(license, (*License).free)
	return license
}

// Get wraps zlm_license_get(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Get(product, version, argv0, path, licenseString string) error {
	var cProduct, cVersion, cArgv0, cPath, cLicenseString *C.char
	// convert argument to C strings
	if product != "" {
		cProduct = C.CString(product)
	}
	if version != "" {
		cVersion = C.CString(version)
	}
	if argv0 != "" {
		cArgv0 = C.CString(argv0)
	}
	if path != "" {
		cPath = C.CString(path)
	}
	if licenseString != "" {
		cLicenseString = C.CString(licenseString)
	}
	// call actual method
	if C.zlm_license_get(license.l, cProduct, cVersion, cArgv0, cPath, cLicenseString, C.zlmgo_errbuf) != C.ZLM_OK {
		return errors.New(C.GoString(C.zlmgo_errbuf))
	}
	return nil
}

func (license *License) free() {
	C.zlm_license_free(license.l)
}

// Product wraps zlm_license_product(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Product() string {
	return C.GoString(C.zlm_license_product(license.l))
}

// Expiry wraps zlm_license_expiry(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Expiry() string {
	return C.GoString(C.zlm_license_expiry(license.l))
}

// ExpiryDays wraps zlm_license_expiry_days(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) ExpiryDays() int {
	return int(C.zlm_license_expiry_days(license.l))
}

// Customer wraps zlm_license_customer(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Customer() string {
	return C.GoString(C.zlm_license_customer(license.l))
}

// Userdata wraps zlm_license_userdata(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Userdata() string {
	return C.GoString(C.zlm_license_userdata(license.l))
}

// UserdataUnescaped wraps zlm_license_userdata_unescaped(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) UserdataUnescaped() string {
	return C.GoString(C.zlm_license_userdata_unescaped(license.l))
}

// Next wraps zlm_license_next(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) Next() error {
	if C.zlm_license_next(license.l, C.zlmgo_errbuf) != C.ZLM_OK {
		return errors.New(C.GoString(C.zlmgo_errbuf))
	}
	return nil
}

// NotExpired wraps zlm_license_not_expired(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) NotExpired() error {
	if C.zlm_license_not_expired(license.l, C.zlmgo_errbuf) != C.ZLM_OK {
		return errors.New(C.GoString(C.zlmgo_errbuf))
	}
	return nil
}

// Version wraps zlm_version(), see
// https://zenlicensemanager.com/documentation/#API
func Version() string {
	return C.GoString(C.zlm_version())
}

// HostidJSON wraps zlm_hostid_json(), see
// https://zenlicensemanager.com/documentation/#API
func HostidJSON() (string, error) {
	cs := C.zlm_hostid_json(C.zlmgo_errbuf)
	if cs == nil {
		return "", errors.New(C.GoString(C.zlmgo_errbuf))
	}
	hostid := C.GoString(cs)
	C.zlm_free(unsafe.Pointer(cs))
	return hostid, nil
}

// CheckA wraps zlm_license_check_a(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckA() {
	C.zlm_license_check_a(license.l)
}

// CheckB wraps zlm_license_check_b(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckB() {
	C.zlm_license_check_b(license.l)
}

// CheckC wraps zlm_license_check_c(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckC() {
	C.zlm_license_check_c(license.l)
}

// CheckD wraps zlm_license_check_d(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckD() {
	C.zlm_license_check_d(license.l)
}

// CheckE wraps zlm_license_check_e(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckE() {
	C.zlm_license_check_e(license.l)
}

// CheckF wraps zlm_license_check_f(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckF() {
	C.zlm_license_check_f(license.l)
}

// CheckG wraps zlm_license_check_g(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckG() {
	C.zlm_license_check_g(license.l)
}

// CheckH wraps zlm_license_check_h(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckH() {
	C.zlm_license_check_h(license.l)
}

// CheckI wraps zlm_license_check_i(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckI() {
	C.zlm_license_check_i(license.l)
}

// CheckJ wraps zlm_license_check_j(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckJ() {
	C.zlm_license_check_j(license.l)
}

// CheckK wraps zlm_license_check_k(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckK() {
	C.zlm_license_check_k(license.l)
}

// CheckL wraps zlm_license_check_l(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckL() {
	C.zlm_license_check_l(license.l)
}

// CheckM wraps zlm_license_check_m(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckM() {
	C.zlm_license_check_m(license.l)
}

// CheckN wraps zlm_license_check_n(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckN() {
	C.zlm_license_check_n(license.l)
}

// CheckO wraps zlm_license_check_o(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckO() {
	C.zlm_license_check_o(license.l)
}

// CheckP wraps zlm_license_check_p(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckP() {
	C.zlm_license_check_p(license.l)
}

// CheckQ wraps zlm_license_check_q(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckQ() {
	C.zlm_license_check_q(license.l)
}

// CheckR wraps zlm_license_check_r(), see
// https://zenlicensemanager.com/documentation/#API
func (license *License) CheckR() {
	C.zlm_license_check_r(license.l)
}
<|endoftext|>


<commit_before>/*
   The MIT License (MIT)

   Copyright (c) 2013 Ondřej Kupka

   Permission is hereby granted, free of charge, to any person obtaining a copy of
   this software and associated documentation files (the "Software"), to deal in
   the Software without restriction, including without limitation the rights to
   use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
   the Software, and to permit persons to whom the Software is furnished to do so,
   subject to the following conditions:

   The above copyright notice and this permission notice shall be included in all
   copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
   FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
   COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
   IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
   CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

package gocli

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"strings"
	"text/template"
)

// Command represents a node in the chain of subcommands. Even App uses it,
// it is embedded there.
type Command struct {
	// Fitting to be set in the Command struct literal.
	UsageLine string
	Short     string
	Long      string
	Action    func(cmd *Command, args []string)

	// Flag set for this subcommand.
	Flags flag.FlagSet

	// All subcommands registered. Although it can be accessed directly,
	// it is not supposed to be used like that. Use MustRegisterSubcmd().
	// This is exported only for the templates to be able to access it.
	Subcmds []*Command

	// Internal field set according to the command type.
	helpTemplate     string
	helpTemplateData interface{}
}

// Get name from the usage line. The name is the first word on the usage line.
func (cmd *Command) Name() string {
	name := strings.TrimSpace(cmd.UsageLine)
	if i := strings.Index(name, " "); i >= 0 {
		name = name[:i]
	}
	return name
}

// Default Command help template.
var CommandHelpTemplate = `COMMAND:
  {{.Name}} - {{.Short}}

USAGE:
  {{.UsageLine}}

OPTIONS:
{{.DefaultFlagsString}}
{{with .Long}}DESCRIPTION:{{.}}{{end}}
{{with .Subcmds}}SUBCOMMANDS:
  {{range .}}{{.Name}}{{with .Short}}{{ "\t" }} - {{.}}{{end}}
  {{end}}
{{end}}`

// Print help and exit.
func (cmd *Command) Usage() {
	t := template.Must(template.New("usage").Parse(cmd.helpTemplate))
	cmd.UsageLine = strings.TrimSpace(cmd.UsageLine)
	t.Execute(os.Stderr, cmd.helpTemplateData)
}

// Returns the same as flag.FlagSet.PrintDefaults, but as a string instead
// of printing it directly into the output stream.
func (cmd *Command) DefaultFlagsString() string {
	var b bytes.Buffer
	cmd.Flags.SetOutput(&b)
	cmd.Flags.PrintDefaults()
	return b.String()
}

// Run the command with supplied arguments. This is called recursively if a
// subcommand is detected, just a prefix is cut off from the arguments.
func (cmd *Command) Run(args []string) {
	err := cmd.Flags.Parse(args)
	if err != nil {
		cmd.Usage()
		os.Exit(1)
	}
	subArgs := cmd.Flags.Args()

	if len(subArgs) == 0 {
		cmd.Action(cmd, subArgs)
		return
	}

	name := subArgs[0]
	for _, subcmd := range cmd.Subcmds {
		if subcmd.Name() == name {
			subcmd.Run(subArgs[1:])
			return
		}
	}

	cmd.Action(cmd, subArgs)
}

// Register a new subcommand with the command. Panic if something is wrong.
func (cmd *Command) MustRegisterSubcommand(subcmd *Command) {
	// Require some fields to be non-empty.
	switch {
	case subcmd.Short == "":
		panic("Short not set")
	case subcmd.UsageLine == "":
		panic("UsageLine not set")
	}

	// Fill in the unexported fields.
	subcmd.helpTemplate = CommandHelpTemplate
	subcmd.helpTemplateData = subcmd

	// Print help if there is no action defined.
	if subcmd.Action == nil {
		subcmd.Action = helpAction(1)
	}

	// Define the help flag.
	subcmd.Flags.Var((*helpValue)(subcmd), "h", "print help and exit")

	// Easy if there are no subcommands defined yet.
	if cmd.Subcmds == nil {
		cmd.Subcmds = []*Command{subcmd}
		return
	}

	// Check for subcommand name collisions.
	for _, c := range cmd.Subcmds {
		if c.Name() == subcmd.Name() {
			panic(fmt.Sprintf("Subcommand %s already defined", subcmd.Name))
		}
	}

	cmd.Subcmds = append(cmd.Subcmds, subcmd)
}
<commit_msg>Use tabwriter to format SUBCOMMANDS section<commit_after>/*
   The MIT License (MIT)

   Copyright (c) 2013 Ondřej Kupka

   Permission is hereby granted, free of charge, to any person obtaining a copy of
   this software and associated documentation files (the "Software"), to deal in
   the Software without restriction, including without limitation the rights to
   use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
   the Software, and to permit persons to whom the Software is furnished to do so,
   subject to the following conditions:

   The above copyright notice and this permission notice shall be included in all
   copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
   FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
   COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
   IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
   CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

package gocli

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"strings"
	"text/tabwriter"
	"text/template"
)

// Command represents a node in the chain of subcommands. Even App uses it,
// it is embedded there.
type Command struct {
	// Fitting to be set in the Command struct literal.
	UsageLine string
	Short     string
	Long      string
	Action    func(cmd *Command, args []string)

	// Flag set for this subcommand.
	Flags flag.FlagSet

	// All subcommands registered. Although it can be accessed directly,
	// it is not supposed to be used like that. Use MustRegisterSubcmd().
	// This is exported only for the templates to be able to access it.
	Subcmds []*Command

	// Internal field set according to the command type.
	helpTemplate     string
	helpTemplateData interface{}
}

// Get name from the usage line. The name is the first word on the usage line.
func (cmd *Command) Name() string {
	name := strings.TrimSpace(cmd.UsageLine)
	if i := strings.Index(name, " "); i >= 0 {
		name = name[:i]
	}
	return name
}

// Default Command help template.
var CommandHelpTemplate = `COMMAND:
  {{.Name}} - {{.Short}}

USAGE:
  {{.UsageLine}}

OPTIONS:
{{.DefaultFlagsString}}
{{with .Long}}DESCRIPTION:{{.}}{{end}}
{{with .Subcmds}}SUBCOMMANDS:
  {{range .}}{{.Name}}{{with .Short}}{{ "\t" }} - {{.}}{{end}}
  {{end}}
{{end}}`

// Print help and exit.
func (cmd *Command) Usage() {
	cmd.UsageLine = strings.TrimSpace(cmd.UsageLine)

	tw := tabwriter.NewWriter(os.Stderr, 0, 8, 2, '\t', 0)
	defer tw.Flush()

	t := template.Must(template.New("usage").Parse(cmd.helpTemplate))
	t.Execute(tw, cmd.helpTemplateData)
}

// Returns the same as flag.FlagSet.PrintDefaults, but as a string instead
// of printing it directly into the output stream.
func (cmd *Command) DefaultFlagsString() string {
	var b bytes.Buffer
	cmd.Flags.SetOutput(&b)
	cmd.Flags.PrintDefaults()
	return b.String()
}

// Run the command with supplied arguments. This is called recursively if a
// subcommand is detected, just a prefix is cut off from the arguments.
func (cmd *Command) Run(args []string) {
	err := cmd.Flags.Parse(args)
	if err != nil {
		cmd.Usage()
		os.Exit(1)
	}
	subArgs := cmd.Flags.Args()

	if len(subArgs) == 0 {
		cmd.Action(cmd, subArgs)
		return
	}

	name := subArgs[0]
	for _, subcmd := range cmd.Subcmds {
		if subcmd.Name() == name {
			subcmd.Run(subArgs[1:])
			return
		}
	}

	cmd.Action(cmd, subArgs)
}

// Register a new subcommand with the command. Panic if something is wrong.
func (cmd *Command) MustRegisterSubcommand(subcmd *Command) {
	// Require some fields to be non-empty.
	switch {
	case subcmd.Short == "":
		panic("Short not set")
	case subcmd.UsageLine == "":
		panic("UsageLine not set")
	}

	// Fill in the unexported fields.
	subcmd.helpTemplate = CommandHelpTemplate
	subcmd.helpTemplateData = subcmd

	// Print help if there is no action defined.
	if subcmd.Action == nil {
		subcmd.Action = helpAction(1)
	}

	// Define the help flag.
	subcmd.Flags.Var((*helpValue)(subcmd), "h", "print help and exit")

	// Easy if there are no subcommands defined yet.
	if cmd.Subcmds == nil {
		cmd.Subcmds = []*Command{subcmd}
		return
	}

	// Check for subcommand name collisions.
	for _, c := range cmd.Subcmds {
		if c.Name() == subcmd.Name() {
			panic(fmt.Sprintf("Subcommand %s already defined", subcmd.Name))
		}
	}

	cmd.Subcmds = append(cmd.Subcmds, subcmd)
}
<|endoftext|>