<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 transports

import (
  "bytes"
  "crypto/tls"
  "crypto/x509"
  "encoding/binary"
  "encoding/pem"
  "errors"
  "fmt"
  "io/ioutil"
  "lc-lib/core"
  "math/rand"
  "net"
  "regexp"
  "sync"
  "time"
)

// Support for newer SSL signature algorithms
import _ "crypto/sha256"
import _ "crypto/sha512"

const (
  // Essentially, this is how often we should check for disconnect/shutdown during socket reads
  socket_interval_seconds = 1
)

type TransportTcpRegistrar struct {
}

type TransportTcpFactory struct {
  transport string

  SSLCertificate string `config:"ssl certificate"`
  SSLKey         string `config:"ssl key"`
  SSLCA          string `config:"ssl ca"`

  hostport_re *regexp.Regexp
  tls_config  tls.Config
}

type TransportTcp struct {
  config     *TransportTcpFactory
  net_config *core.NetworkConfig
  socket     net.Conn
  tlssocket  *tls.Conn

  wait     sync.WaitGroup
  shutdown chan interface{}

  send_chan chan []byte
  recv_chan chan interface{}

  can_send chan int
}

func NewTcpTransportFactory(config *core.Config, config_path string, unused map[string]interface{}, name string) (core.TransportFactory, error) {
  var err error

  ret := &TransportTcpFactory{
    transport:   name,
    hostport_re: regexp.MustCompile(`^\[?([^]]+)\]?:([0-9]+)$`),
  }

  // Only allow SSL configurations if this is "tls"
  if name == "tls" {
    if err = config.PopulateConfig(ret, config_path, unused); err != nil {
      return nil, err
    }

    if len(ret.SSLCertificate) > 0 && len(ret.SSLKey) > 0 {
      cert, err := tls.LoadX509KeyPair(ret.SSLCertificate, ret.SSLKey)
      if err != nil {
        return nil, fmt.Errorf("Failed loading client ssl certificate: %s", err)
      }

      ret.tls_config.Certificates = []tls.Certificate{cert}
    }

    if len(ret.SSLCA) > 0 {
      ret.tls_config.RootCAs = x509.NewCertPool()

      pemdata, err := ioutil.ReadFile(ret.SSLCA)
      if err != nil {
        return nil, fmt.Errorf("Failure reading CA certificate: %s", err)
      }

      block, _ := pem.Decode(pemdata)
      if block == nil {
        return nil, errors.New("Failed to decode CA certificate data")
      }
      if block.Type != "CERTIFICATE" {
        return nil, fmt.Errorf("Specified CA certificate is not a certificate: %s", ret.SSLCA)
      }

      cert, err := x509.ParseCertificate(block.Bytes)
      if err != nil {
        return nil, fmt.Errorf("Failed to parse CA certificate: %s", err)
      }

      ret.tls_config.RootCAs.AddCert(cert)
    }
  } else {
    if err := config.ReportUnusedConfig(config_path, unused); err != nil {
      return nil, err
    }
  }

  return ret, nil
}

func (f *TransportTcpFactory) NewTransport(config *core.NetworkConfig) (core.Transport, error) {
  return &TransportTcp{config: f, net_config: config}, nil
}

func (t *TransportTcp) ReloadConfig(new_net_config *core.NetworkConfig) int {
  // Check we can grab new TCP config to compare, if not force transport reinit
  new_config, ok := new_net_config.TransportFactory.(*TransportTcpFactory)
  if !ok {
    return core.Reload_Transport
  }

  if new_config.SSLCertificate != t.config.SSLCertificate || new_config.SSLKey != t.config.SSLKey || new_config.SSLCA != t.config.SSLCA {
    return core.Reload_Transport
  }

  // Publisher handles changes to net_config, but ensure we store the latest in case it asks for a reconnect
  t.net_config = new_net_config

  return core.Reload_None
}

func (t *TransportTcp) Init() error {
  if t.shutdown != nil {
    t.disconnect()
  }

  // Pick a random server from the list.
  hostport := t.net_config.Servers[rand.Int()%len(t.net_config.Servers)]
  submatch := t.config.hostport_re.FindSubmatch([]byte(hostport))
  if submatch == nil {
    return fmt.Errorf("Invalid host:port given: %s", hostport)
  }

  // Lookup the server in DNS (if this is IP it will implicitly return)
  host := string(submatch[1])
  port := string(submatch[2])
  addresses, err := net.LookupHost(host)
  if err != nil {
    return fmt.Errorf("DNS lookup failure \"%s\": %s", host, err)
  }

  // Select a random address from the pool of addresses provided by DNS
  address := addresses[rand.Int()%len(addresses)]
  addressport := net.JoinHostPort(address, port)

  log.Info("Attempting to connect to %s (%s)", addressport, host)

  tcpsocket, err := net.DialTimeout("tcp", addressport, t.net_config.Timeout)
  if err != nil {
    return fmt.Errorf("Failed to connect to %s: %s", address, err)
  }

  // Now wrap in TLS if this is the "tls" transport
  if t.config.transport == "tls" {
    // Set the tlsconfig server name for server validation (since Go 1.3)
    t.config.tls_config.ServerName = host

    t.tlssocket = tls.Client(&transportTcpWrap{transport: t, tcpsocket: tcpsocket}, &t.config.tls_config)
    t.tlssocket.SetDeadline(time.Now().Add(t.net_config.Timeout))
    err = t.tlssocket.Handshake()
    if err != nil {
      t.tlssocket.Close()
      tcpsocket.Close()
      return fmt.Errorf("TLS Handshake failure with %s: %s", address, err)
    }

    t.socket = t.tlssocket
  } else {
    t.socket = tcpsocket
  }

  log.Info("Connected to %s", address)

  // Signal channels
  t.shutdown = make(chan interface{}, 1)
  t.send_chan = make(chan []byte, 1)
  // Buffer of two for recv_chan since both routines may send an error to it
  // First error we get back initiates disconnect, thus we must not block routines
  t.recv_chan = make(chan interface{}, 2)
  t.can_send = make(chan int, 1)

  // Start with a send
  t.can_send <- 1

  t.wait.Add(2)

  // Start separate sender and receiver so we can asynchronously send and receive for max performance
  // They have to be different routines too because we don't have cross-platform poll, so they will need to block
  // Of course, we'll time out and check shutdown on occasion
  go t.sender()
  go t.receiver()

  return nil
}

func (t *TransportTcp) disconnect() {
  if t.shutdown == nil {
    return
  }

  // Send shutdown request
  close(t.shutdown)
  t.wait.Wait()
  t.shutdown = nil

  // If tls, shutdown tls socket first
  if t.config.transport == "tls" {
    t.tlssocket.Close()
  }

  t.socket.Close()
}

func (t *TransportTcp) sender() {
SendLoop:
  for {
    select {
    case <-t.shutdown:
      // Shutdown
      break SendLoop
    case msg := <-t.send_chan:
      // Ask for more while we send this
      t.setChan(t.can_send)
      // Write deadline is managed by our net.Conn wrapper that tls will call into
      _, err := t.socket.Write(msg)
      if err != nil {
        if net_err, ok := err.(net.Error); ok && net_err.Timeout() {
          // Shutdown will have been received by the wrapper
          break SendLoop
        } else {
          // Pass error back
          t.recv_chan <- err
        }
      }
    }
  }

  t.wait.Done()
}

func (t *TransportTcp) receiver() {
  var err error
  var shutdown bool
  header := make([]byte, 8)

  for {
    if err, shutdown = t.receiverRead(header); err != nil || shutdown {
      break
    }

    // Grab length of message
    length := binary.BigEndian.Uint32(header[4:8])

    // Sanity
    if length > 1048576 {
      err = fmt.Errorf("Data too large (%d)", length)
      break
    }

    // Allocate for full message
    message := make([]byte, length)

    if err, shutdown = t.receiverRead(message); err != nil || shutdown {
      break
    }

    // Pass back the message
    t.recv_chan <- [][]byte{header[0:4], message}
  } /* loop until shutdown */

  if err != nil {
    // Pass the error back and abort
    t.recv_chan <- err
  }

  t.wait.Done()
}

func (t *TransportTcp) receiverRead(data []byte) (error, bool) {
  received := 0

RecvLoop:
  for {
    select {
    case <-t.shutdown:
      // Shutdown
      break RecvLoop
    default:
      // Timeout after socket_interval_seconds, check for shutdown, and try again
      t.socket.SetReadDeadline(time.Now().Add(socket_interval_seconds * time.Second))

      length, err := t.socket.Read(data[received:])
      received += length
      if err == nil || received >= len(data) {
        // Success
        return nil, false
      } else if net_err, ok := err.(net.Error); ok && net_err.Timeout() {
        // Keep trying
        continue
      } else {
        // Pass an error back
        return err, false
      }
    } /* select */
  } /* loop until required amount receive or shutdown */

  return nil, true
}

func (t *TransportTcp) setChan(set chan<- int) {
  select {
  case set <- 1:
  default:
  }
}

func (t *TransportTcp) CanSend() <-chan int {
  return t.can_send
}

func (t *TransportTcp) Write(signature string, message []byte) (err error) {
  var write_buffer *bytes.Buffer
  write_buffer = bytes.NewBuffer(make([]byte, 0, len(signature)+4+len(message)))

  if _, err = write_buffer.Write([]byte(signature)); err != nil {
    return
  }
  if err = binary.Write(write_buffer, binary.BigEndian, uint32(len(message))); err != nil {
    return
  }
  if len(message) != 0 {
    if _, err = write_buffer.Write(message); err != nil {
      return
    }
  }

  t.send_chan <- write_buffer.Bytes()
  return nil
}

func (t *TransportTcp) Read() <-chan interface{} {
  return t.recv_chan
}

func (t *TransportTcp) Shutdown() {
  t.disconnect()
}

// Register the transports
func init() {
  rand.Seed(time.Now().UnixNano())

  core.RegisterTransport("tcp", NewTcpTransportFactory)
  core.RegisterTransport("tls", NewTcpTransportFactory)
}
<commit_msg>Fix random hang in TCP reconnect in situations where the connection succeeds and immediately fails due to a certificate problem etc<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 transports

import (
  "bytes"
  "crypto/tls"
  "crypto/x509"
  "encoding/binary"
  "encoding/pem"
  "errors"
  "fmt"
  "io/ioutil"
  "lc-lib/core"
  "math/rand"
  "net"
  "regexp"
  "sync"
  "time"
)

// Support for newer SSL signature algorithms
import _ "crypto/sha256"
import _ "crypto/sha512"

const (
  // Essentially, this is how often we should check for disconnect/shutdown during socket reads
  socket_interval_seconds = 1
)

type TransportTcpRegistrar struct {
}

type TransportTcpFactory struct {
  transport string

  SSLCertificate string `config:"ssl certificate"`
  SSLKey         string `config:"ssl key"`
  SSLCA          string `config:"ssl ca"`

  hostport_re *regexp.Regexp
  tls_config  tls.Config
}

type TransportTcp struct {
  config     *TransportTcpFactory
  net_config *core.NetworkConfig
  socket     net.Conn
  tlssocket  *tls.Conn

  wait     sync.WaitGroup
  shutdown chan interface{}

  send_chan chan []byte
  recv_chan chan interface{}

  can_send chan int
}

func NewTcpTransportFactory(config *core.Config, config_path string, unused map[string]interface{}, name string) (core.TransportFactory, error) {
  var err error

  ret := &TransportTcpFactory{
    transport:   name,
    hostport_re: regexp.MustCompile(`^\[?([^]]+)\]?:([0-9]+)$`),
  }

  // Only allow SSL configurations if this is "tls"
  if name == "tls" {
    if err = config.PopulateConfig(ret, config_path, unused); err != nil {
      return nil, err
    }

    if len(ret.SSLCertificate) > 0 && len(ret.SSLKey) > 0 {
      cert, err := tls.LoadX509KeyPair(ret.SSLCertificate, ret.SSLKey)
      if err != nil {
        return nil, fmt.Errorf("Failed loading client ssl certificate: %s", err)
      }

      ret.tls_config.Certificates = []tls.Certificate{cert}
    }

    if len(ret.SSLCA) > 0 {
      ret.tls_config.RootCAs = x509.NewCertPool()

      pemdata, err := ioutil.ReadFile(ret.SSLCA)
      if err != nil {
        return nil, fmt.Errorf("Failure reading CA certificate: %s", err)
      }

      block, _ := pem.Decode(pemdata)
      if block == nil {
        return nil, errors.New("Failed to decode CA certificate data")
      }
      if block.Type != "CERTIFICATE" {
        return nil, fmt.Errorf("Specified CA certificate is not a certificate: %s", ret.SSLCA)
      }

      cert, err := x509.ParseCertificate(block.Bytes)
      if err != nil {
        return nil, fmt.Errorf("Failed to parse CA certificate: %s", err)
      }

      ret.tls_config.RootCAs.AddCert(cert)
    }
  } else {
    if err := config.ReportUnusedConfig(config_path, unused); err != nil {
      return nil, err
    }
  }

  return ret, nil
}

func (f *TransportTcpFactory) NewTransport(config *core.NetworkConfig) (core.Transport, error) {
  return &TransportTcp{config: f, net_config: config}, nil
}

func (t *TransportTcp) ReloadConfig(new_net_config *core.NetworkConfig) int {
  // Check we can grab new TCP config to compare, if not force transport reinit
  new_config, ok := new_net_config.TransportFactory.(*TransportTcpFactory)
  if !ok {
    return core.Reload_Transport
  }

  if new_config.SSLCertificate != t.config.SSLCertificate || new_config.SSLKey != t.config.SSLKey || new_config.SSLCA != t.config.SSLCA {
    return core.Reload_Transport
  }

  // Publisher handles changes to net_config, but ensure we store the latest in case it asks for a reconnect
  t.net_config = new_net_config

  return core.Reload_None
}

func (t *TransportTcp) Init() error {
  if t.shutdown != nil {
    t.disconnect()
  }

  // Pick a random server from the list.
  hostport := t.net_config.Servers[rand.Int()%len(t.net_config.Servers)]
  submatch := t.config.hostport_re.FindSubmatch([]byte(hostport))
  if submatch == nil {
    return fmt.Errorf("Invalid host:port given: %s", hostport)
  }

  // Lookup the server in DNS (if this is IP it will implicitly return)
  host := string(submatch[1])
  port := string(submatch[2])
  addresses, err := net.LookupHost(host)
  if err != nil {
    return fmt.Errorf("DNS lookup failure \"%s\": %s", host, err)
  }

  // Select a random address from the pool of addresses provided by DNS
  address := addresses[rand.Int()%len(addresses)]
  addressport := net.JoinHostPort(address, port)

  log.Info("Attempting to connect to %s (%s)", addressport, host)

  tcpsocket, err := net.DialTimeout("tcp", addressport, t.net_config.Timeout)
  if err != nil {
    return fmt.Errorf("Failed to connect to %s: %s", address, err)
  }

  // Now wrap in TLS if this is the "tls" transport
  if t.config.transport == "tls" {
    // Set the tlsconfig server name for server validation (since Go 1.3)
    t.config.tls_config.ServerName = host

    t.tlssocket = tls.Client(&transportTcpWrap{transport: t, tcpsocket: tcpsocket}, &t.config.tls_config)
    t.tlssocket.SetDeadline(time.Now().Add(t.net_config.Timeout))
    err = t.tlssocket.Handshake()
    if err != nil {
      t.tlssocket.Close()
      tcpsocket.Close()
      return fmt.Errorf("TLS Handshake failure with %s: %s", address, err)
    }

    t.socket = t.tlssocket
  } else {
    t.socket = tcpsocket
  }

  log.Info("Connected to %s", address)

  // Signal channels
  t.shutdown = make(chan interface{}, 1)
  t.send_chan = make(chan []byte, 1)
  // Buffer of two for recv_chan since both routines may send an error to it
  // First error we get back initiates disconnect, thus we must not block routines
  t.recv_chan = make(chan interface{}, 2)
  t.can_send = make(chan int, 1)

  // Start with a send
  t.can_send <- 1

  t.wait.Add(2)

  // Start separate sender and receiver so we can asynchronously send and receive for max performance
  // They have to be different routines too because we don't have cross-platform poll, so they will need to block
  // Of course, we'll time out and check shutdown on occasion
  go t.sender()
  go t.receiver()

  return nil
}

func (t *TransportTcp) disconnect() {
  if t.shutdown == nil {
    return
  }

  // Send shutdown request
  close(t.shutdown)
  t.wait.Wait()
  t.shutdown = nil

  // If tls, shutdown tls socket first
  if t.config.transport == "tls" {
    t.tlssocket.Close()
  }

  t.socket.Close()
}

func (t *TransportTcp) sender() {
SendLoop:
  for {
    select {
    case <-t.shutdown:
      // Shutdown
      break SendLoop
    case msg := <-t.send_chan:
      // Ask for more while we send this
      t.setChan(t.can_send)
      // Write deadline is managed by our net.Conn wrapper that tls will call into
      _, err := t.socket.Write(msg)
      if err != nil {
        if net_err, ok := err.(net.Error); ok && net_err.Timeout() {
          // Shutdown will have been received by the wrapper
          break SendLoop
        } else {
          // Pass error back
          t.recv_chan <- err
        }
      }
    }
  }

  t.wait.Done()
}

func (t *TransportTcp) receiver() {
  var err error
  var shutdown bool
  header := make([]byte, 8)

  for {
    if err, shutdown = t.receiverRead(header); err != nil || shutdown {
      break
    }

    // Grab length of message
    length := binary.BigEndian.Uint32(header[4:8])

    // Sanity
    if length > 1048576 {
      err = fmt.Errorf("Data too large (%d)", length)
      break
    }

    // Allocate for full message
    message := make([]byte, length)

    if err, shutdown = t.receiverRead(message); err != nil || shutdown {
      break
    }

    // Pass back the message
    select {
    case <-t.shutdown:
      break
    case t.recv_chan <- [][]byte{header[0:4], message}:
    }
  } /* loop until shutdown */

  if err != nil {
    // Pass the error back and abort
    select {
    case <-t.shutdown:
    case t.recv_chan <- err:
    }
  }

  t.wait.Done()
}

func (t *TransportTcp) receiverRead(data []byte) (error, bool) {
  received := 0

RecvLoop:
  for {
    select {
    case <-t.shutdown:
      // Shutdown
      break RecvLoop
    default:
      // Timeout after socket_interval_seconds, check for shutdown, and try again
      t.socket.SetReadDeadline(time.Now().Add(socket_interval_seconds * time.Second))

      length, err := t.socket.Read(data[received:])
      received += length
      if err == nil || received >= len(data) {
        // Success
        return nil, false
      } else if net_err, ok := err.(net.Error); ok && net_err.Timeout() {
        // Keep trying
        continue
      } else {
        // Pass an error back
        return err, false
      }
    } /* select */
  } /* loop until required amount receive or shutdown */

  return nil, true
}

func (t *TransportTcp) setChan(set chan<- int) {
  select {
  case set <- 1:
  default:
  }
}

func (t *TransportTcp) CanSend() <-chan int {
  return t.can_send
}

func (t *TransportTcp) Write(signature string, message []byte) (err error) {
  var write_buffer *bytes.Buffer
  write_buffer = bytes.NewBuffer(make([]byte, 0, len(signature)+4+len(message)))

  if _, err = write_buffer.Write([]byte(signature)); err != nil {
    return
  }
  if err = binary.Write(write_buffer, binary.BigEndian, uint32(len(message))); err != nil {
    return
  }
  if len(message) != 0 {
    if _, err = write_buffer.Write(message); err != nil {
      return
    }
  }

  t.send_chan <- write_buffer.Bytes()
  return nil
}

func (t *TransportTcp) Read() <-chan interface{} {
  return t.recv_chan
}

func (t *TransportTcp) Shutdown() {
  t.disconnect()
}

// Register the transports
func init() {
  rand.Seed(time.Now().UnixNano())

  core.RegisterTransport("tcp", NewTcpTransportFactory)
  core.RegisterTransport("tls", NewTcpTransportFactory)
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ssm"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsSsmDocument() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSsmDocumentCreate,
		Read:   resourceAwsSsmDocumentRead,
		Update: resourceAwsSsmDocumentUpdate,
		Delete: resourceAwsSsmDocumentDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
			},
			"content": {
				Type:     schema.TypeString,
				Required: true,
			},
			"document_type": {
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validateAwsSSMDocumentType,
			},
			"created_date": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"default_version": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"description": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"hash": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"hash_type": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"latest_version": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"owner": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"status": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"platform_types": {
				Type:     schema.TypeList,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},
			"parameter": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"default_value": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"description": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"type": {
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},
			"permissions": {
				Type:     schema.TypeMap,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"type": {
							Type:     schema.TypeString,
							Required: true,
						},
						"account_ids": {
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
			},
		},
	}
}

func resourceAwsSsmDocumentCreate(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Creating SSM Document: %s", d.Get("name").(string))

	docInput := &ssm.CreateDocumentInput{
		Name:         aws.String(d.Get("name").(string)),
		Content:      aws.String(d.Get("content").(string)),
		DocumentType: aws.String(d.Get("document_type").(string)),
	}

	log.Printf("[DEBUG] Waiting for SSM Document %q to be created", d.Get("name").(string))
	err := resource.Retry(5*time.Minute, func() *resource.RetryError {
		resp, err := ssmconn.CreateDocument(docInput)

		if err != nil {
			return resource.NonRetryableError(err)
		}

		d.SetId(*resp.DocumentDescription.Name)
		return nil
	})

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error creating SSM document: {{err}}", err)
	}

	if v, ok := d.GetOk("permissions"); ok && v != nil {
		if err := setDocumentPermissions(d, meta); err != nil {
			return err
		}
	} else {
		log.Printf("[DEBUG] Not setting permissions for %q", d.Id())
	}

	return resourceAwsSsmDocumentRead(d, meta)
}

func resourceAwsSsmDocumentRead(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[DEBUG] Reading SSM Document: %s", d.Id())

	docInput := &ssm.DescribeDocumentInput{
		Name: aws.String(d.Get("name").(string)),
	}

	resp, err := ssmconn.DescribeDocument(docInput)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error describing SSM document: {{err}}", err)
	}

	doc := resp.Document
	d.Set("created_date", doc.CreatedDate)
	d.Set("default_version", doc.DefaultVersion)
	d.Set("description", doc.Description)

	if _, ok := d.GetOk("document_type"); ok {
		d.Set("document_type", doc.DocumentType)
	}

	d.Set("document_version", doc.DocumentVersion)
	d.Set("hash", doc.Hash)
	d.Set("hash_type", doc.HashType)
	d.Set("latest_version", doc.LatestVersion)
	d.Set("name", doc.Name)
	d.Set("owner", doc.Owner)
	d.Set("platform_types", flattenStringList(doc.PlatformTypes))

	d.Set("status", doc.Status)

	gp, err := getDocumentPermissions(d, meta)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error reading SSM document permissions: {{err}}", err)
	}

	d.Set("permissions", gp)

	params := make([]map[string]interface{}, 0)
	for i := 0; i < len(doc.Parameters); i++ {

		dp := doc.Parameters[i]
		param := make(map[string]interface{})

		if dp.DefaultValue != nil {
			param["default_value"] = *dp.DefaultValue
		}
		param["description"] = *dp.Description
		param["name"] = *dp.Name
		param["type"] = *dp.Type
		params = append(params, param)
	}

	if len(params) == 0 {
		params = make([]map[string]interface{}, 1)
	}

	if err := d.Set("parameter", params); err != nil {
		return err
	}

	return nil
}

func resourceAwsSsmDocumentUpdate(d *schema.ResourceData, meta interface{}) error {

	if _, ok := d.GetOk("permissions"); ok {
		if err := setDocumentPermissions(d, meta); err != nil {
			return err
		}
	} else {
		log.Printf("[DEBUG] Not setting document permissions on %q", d.Id())
	}

	return resourceAwsSsmDocumentRead(d, meta)
}

func resourceAwsSsmDocumentDelete(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	if err := deleteDocumentPermissions(d, meta); err != nil {
		return err
	}

	log.Printf("[INFO] Deleting SSM Document: %s", d.Id())

	params := &ssm.DeleteDocumentInput{
		Name: aws.String(d.Get("name").(string)),
	}

	_, err := ssmconn.DeleteDocument(params)
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Waiting for SSM Document %q to be deleted", d.Get("name").(string))
	err = resource.Retry(10*time.Minute, func() *resource.RetryError {
		_, err := ssmconn.DescribeDocument(&ssm.DescribeDocumentInput{
			Name: aws.String(d.Get("name").(string)),
		})

		if err != nil {
			awsErr, ok := err.(awserr.Error)
			if !ok {
				return resource.NonRetryableError(err)
			}

			if awsErr.Code() == "InvalidDocument" {
				return nil
			}

			return resource.NonRetryableError(err)
		}

		return resource.RetryableError(
			fmt.Errorf("%q: Timeout while waiting for the document to be deleted", d.Id()))
	})
	if err != nil {
		return err
	}

	d.SetId("")

	return nil
}

func setDocumentPermissions(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Setting permissions for document: %s", d.Id())
	permission := d.Get("permissions").(map[string]interface{})

	ids := aws.StringSlice([]string{permission["account_ids"].(string)})

	if strings.Contains(permission["account_ids"].(string), ",") {
		ids = aws.StringSlice(strings.Split(permission["account_ids"].(string), ","))
	}

	permInput := &ssm.ModifyDocumentPermissionInput{
		Name:            aws.String(d.Get("name").(string)),
		PermissionType:  aws.String(permission["type"].(string)),
		AccountIdsToAdd: ids,
	}

	_, err := ssmconn.ModifyDocumentPermission(permInput)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error setting permissions for SSM document: {{err}}", err)
	}

	return nil
}

func getDocumentPermissions(d *schema.ResourceData, meta interface{}) (map[string]interface{}, error) {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Getting permissions for document: %s", d.Id())

	//How to get from nested scheme resource?
	permissionType := "Share"

	permInput := &ssm.DescribeDocumentPermissionInput{
		Name:           aws.String(d.Get("name").(string)),
		PermissionType: aws.String(permissionType),
	}

	resp, err := ssmconn.DescribeDocumentPermission(permInput)

	if err != nil {
		return nil, errwrap.Wrapf("[ERROR] Error setting permissions for SSM document: {{err}}", err)
	}

	var account_ids = make([]string, len(resp.AccountIds))
	for i := 0; i < len(resp.AccountIds); i++ {
		account_ids[i] = *resp.AccountIds[i]
	}

	var ids = ""
	if len(account_ids) == 1 {
		ids = account_ids[0]
	} else if len(account_ids) > 1 {
		ids = strings.Join(account_ids, ",")
	} else {
		ids = ""
	}

	if ids == "" {
		return nil, nil
	}

	perms := make(map[string]interface{})
	perms["type"] = permissionType
	perms["account_ids"] = ids

	return perms, nil
}

func deleteDocumentPermissions(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Removing permissions from document: %s", d.Id())

	permInput := &ssm.ModifyDocumentPermissionInput{
		Name:               aws.String(d.Get("name").(string)),
		PermissionType:     aws.String("Share"),
		AccountIdsToRemove: aws.StringSlice(strings.Split("all", ",")),
	}

	_, err := ssmconn.ModifyDocumentPermission(permInput)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error removing permissions for SSM document: {{err}}", err)
	}

	return nil
}

func validateAwsSSMDocumentType(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	types := map[string]bool{
		"Command":    true,
		"Policy":     true,
		"Automation": true,
	}

	if !types[value] {
		errors = append(errors, fmt.Errorf("CodeBuild: Arifacts Namespace Type can only be NONE / BUILD_ID"))
	}
	return
}
<commit_msg>provider/aws: Fix error message returned when an invalid SSM document type is provided. The original error appears to have been copied from the CodeBuild resource and provides incorrect information about valid types.<commit_after>package aws

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ssm"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsSsmDocument() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsSsmDocumentCreate,
		Read:   resourceAwsSsmDocumentRead,
		Update: resourceAwsSsmDocumentUpdate,
		Delete: resourceAwsSsmDocumentDelete,

		Schema: map[string]*schema.Schema{
			"name": {
				Type:     schema.TypeString,
				Required: true,
			},
			"content": {
				Type:     schema.TypeString,
				Required: true,
			},
			"document_type": {
				Type:         schema.TypeString,
				Required:     true,
				ValidateFunc: validateAwsSSMDocumentType,
			},
			"created_date": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"default_version": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"description": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"hash": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"hash_type": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"latest_version": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"owner": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"status": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"platform_types": {
				Type:     schema.TypeList,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},
			"parameter": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"name": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"default_value": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"description": {
							Type:     schema.TypeString,
							Optional: true,
						},
						"type": {
							Type:     schema.TypeString,
							Optional: true,
						},
					},
				},
			},
			"permissions": {
				Type:     schema.TypeMap,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"type": {
							Type:     schema.TypeString,
							Required: true,
						},
						"account_ids": {
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
			},
		},
	}
}

func resourceAwsSsmDocumentCreate(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Creating SSM Document: %s", d.Get("name").(string))

	docInput := &ssm.CreateDocumentInput{
		Name:         aws.String(d.Get("name").(string)),
		Content:      aws.String(d.Get("content").(string)),
		DocumentType: aws.String(d.Get("document_type").(string)),
	}

	log.Printf("[DEBUG] Waiting for SSM Document %q to be created", d.Get("name").(string))
	err := resource.Retry(5*time.Minute, func() *resource.RetryError {
		resp, err := ssmconn.CreateDocument(docInput)

		if err != nil {
			return resource.NonRetryableError(err)
		}

		d.SetId(*resp.DocumentDescription.Name)
		return nil
	})

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error creating SSM document: {{err}}", err)
	}

	if v, ok := d.GetOk("permissions"); ok && v != nil {
		if err := setDocumentPermissions(d, meta); err != nil {
			return err
		}
	} else {
		log.Printf("[DEBUG] Not setting permissions for %q", d.Id())
	}

	return resourceAwsSsmDocumentRead(d, meta)
}

func resourceAwsSsmDocumentRead(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[DEBUG] Reading SSM Document: %s", d.Id())

	docInput := &ssm.DescribeDocumentInput{
		Name: aws.String(d.Get("name").(string)),
	}

	resp, err := ssmconn.DescribeDocument(docInput)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error describing SSM document: {{err}}", err)
	}

	doc := resp.Document
	d.Set("created_date", doc.CreatedDate)
	d.Set("default_version", doc.DefaultVersion)
	d.Set("description", doc.Description)

	if _, ok := d.GetOk("document_type"); ok {
		d.Set("document_type", doc.DocumentType)
	}

	d.Set("document_version", doc.DocumentVersion)
	d.Set("hash", doc.Hash)
	d.Set("hash_type", doc.HashType)
	d.Set("latest_version", doc.LatestVersion)
	d.Set("name", doc.Name)
	d.Set("owner", doc.Owner)
	d.Set("platform_types", flattenStringList(doc.PlatformTypes))

	d.Set("status", doc.Status)

	gp, err := getDocumentPermissions(d, meta)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error reading SSM document permissions: {{err}}", err)
	}

	d.Set("permissions", gp)

	params := make([]map[string]interface{}, 0)
	for i := 0; i < len(doc.Parameters); i++ {

		dp := doc.Parameters[i]
		param := make(map[string]interface{})

		if dp.DefaultValue != nil {
			param["default_value"] = *dp.DefaultValue
		}
		param["description"] = *dp.Description
		param["name"] = *dp.Name
		param["type"] = *dp.Type
		params = append(params, param)
	}

	if len(params) == 0 {
		params = make([]map[string]interface{}, 1)
	}

	if err := d.Set("parameter", params); err != nil {
		return err
	}

	return nil
}

func resourceAwsSsmDocumentUpdate(d *schema.ResourceData, meta interface{}) error {

	if _, ok := d.GetOk("permissions"); ok {
		if err := setDocumentPermissions(d, meta); err != nil {
			return err
		}
	} else {
		log.Printf("[DEBUG] Not setting document permissions on %q", d.Id())
	}

	return resourceAwsSsmDocumentRead(d, meta)
}

func resourceAwsSsmDocumentDelete(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	if err := deleteDocumentPermissions(d, meta); err != nil {
		return err
	}

	log.Printf("[INFO] Deleting SSM Document: %s", d.Id())

	params := &ssm.DeleteDocumentInput{
		Name: aws.String(d.Get("name").(string)),
	}

	_, err := ssmconn.DeleteDocument(params)
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Waiting for SSM Document %q to be deleted", d.Get("name").(string))
	err = resource.Retry(10*time.Minute, func() *resource.RetryError {
		_, err := ssmconn.DescribeDocument(&ssm.DescribeDocumentInput{
			Name: aws.String(d.Get("name").(string)),
		})

		if err != nil {
			awsErr, ok := err.(awserr.Error)
			if !ok {
				return resource.NonRetryableError(err)
			}

			if awsErr.Code() == "InvalidDocument" {
				return nil
			}

			return resource.NonRetryableError(err)
		}

		return resource.RetryableError(
			fmt.Errorf("%q: Timeout while waiting for the document to be deleted", d.Id()))
	})
	if err != nil {
		return err
	}

	d.SetId("")

	return nil
}

func setDocumentPermissions(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Setting permissions for document: %s", d.Id())
	permission := d.Get("permissions").(map[string]interface{})

	ids := aws.StringSlice([]string{permission["account_ids"].(string)})

	if strings.Contains(permission["account_ids"].(string), ",") {
		ids = aws.StringSlice(strings.Split(permission["account_ids"].(string), ","))
	}

	permInput := &ssm.ModifyDocumentPermissionInput{
		Name:            aws.String(d.Get("name").(string)),
		PermissionType:  aws.String(permission["type"].(string)),
		AccountIdsToAdd: ids,
	}

	_, err := ssmconn.ModifyDocumentPermission(permInput)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error setting permissions for SSM document: {{err}}", err)
	}

	return nil
}

func getDocumentPermissions(d *schema.ResourceData, meta interface{}) (map[string]interface{}, error) {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Getting permissions for document: %s", d.Id())

	//How to get from nested scheme resource?
	permissionType := "Share"

	permInput := &ssm.DescribeDocumentPermissionInput{
		Name:           aws.String(d.Get("name").(string)),
		PermissionType: aws.String(permissionType),
	}

	resp, err := ssmconn.DescribeDocumentPermission(permInput)

	if err != nil {
		return nil, errwrap.Wrapf("[ERROR] Error setting permissions for SSM document: {{err}}", err)
	}

	var account_ids = make([]string, len(resp.AccountIds))
	for i := 0; i < len(resp.AccountIds); i++ {
		account_ids[i] = *resp.AccountIds[i]
	}

	var ids = ""
	if len(account_ids) == 1 {
		ids = account_ids[0]
	} else if len(account_ids) > 1 {
		ids = strings.Join(account_ids, ",")
	} else {
		ids = ""
	}

	if ids == "" {
		return nil, nil
	}

	perms := make(map[string]interface{})
	perms["type"] = permissionType
	perms["account_ids"] = ids

	return perms, nil
}

func deleteDocumentPermissions(d *schema.ResourceData, meta interface{}) error {
	ssmconn := meta.(*AWSClient).ssmconn

	log.Printf("[INFO] Removing permissions from document: %s", d.Id())

	permInput := &ssm.ModifyDocumentPermissionInput{
		Name:               aws.String(d.Get("name").(string)),
		PermissionType:     aws.String("Share"),
		AccountIdsToRemove: aws.StringSlice(strings.Split("all", ",")),
	}

	_, err := ssmconn.ModifyDocumentPermission(permInput)

	if err != nil {
		return errwrap.Wrapf("[ERROR] Error removing permissions for SSM document: {{err}}", err)
	}

	return nil
}

func validateAwsSSMDocumentType(v interface{}, k string) (ws []string, errors []error) {
	value := v.(string)
	types := map[string]bool{
		"Command":    true,
		"Policy":     true,
		"Automation": true,
	}

	if !types[value] {
		errors = append(errors, fmt.Errorf("Document type %s is invalid. Valid types are Command, Policy or Automation", value))
	}
	return
}
<|endoftext|>


<commit_before>package onkyo

import "github.com/cnf/go-claw/tools"
import "strconv"
import "net"
import "fmt"
import "time"

const ONKYO_PORT = 60128
const ONKYO_MAGIC = "!xECNQSTN"

func runOnkyoDetect(ch chan string, timeout int) {
    defer close(ch)
    addrs := tools.BroadcastAddrs()
    port := strconv.Itoa(ONKYO_PORT)

    cmd := &OnkyoCommandTCP{ONKYO_MAGIC}
    transmit_str, err := cmd.Bytes()

    if err != nil {
        return
    }

    udpaddr, err := net.ResolveUDPAddr("udp4", ":0")
    if err != nil {
        return
    }
    conn, err := net.ListenUDP("udp4", udpaddr)
    if err != nil {
        fmt.Printf("ERROR: %s\n", err.Error())
        return
    }
    defer conn.Close()
    conn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Millisecond))

    // Send broadcast
    for _, a := range addrs {
        udpdest, err := net.ResolveUDPAddr("udp4", a+":"+port)
        if err != nil {
            continue
        }
        conn.WriteToUDP(transmit_str, udpdest)
        // Ignore errors - just continue
    }

    // Now loop listening to responses with timeout
    buf := make([]byte, 1024)
    for {
        rlen, raddr, err := conn.ReadFromUDP(buf);
        if err != nil {
            if nerr, ok := err.(net.Error); !ok || !nerr.Temporary() || nerr.Timeout() {
                return
            }
            continue
        }
        //fmt.Printf("Got response from: %v (len: %d):\n", raddr, rlen)
        resp := &OnkyoCommandTCP{}
        if err := resp.Parse(buf[0:rlen]); err != nil {
            continue
        }
        ch <- raddr.String() + ":" + resp.Message()

        // Once we got a response, wait maximum 50ms for other responses
        conn.SetReadDeadline(time.Now().Add(time.Duration(50) * time.Millisecond))
    }
}

func OnkyoAutoDetect(timeout int) []string {
    ch := make(chan string)
    go runOnkyoDetect(ch, timeout)
    ret := make([]string, 0)
    //timer := time.NewTimer(time.Millisecond * time.Duration(timeout + 200))
    for {
        select {
            case s, ok := <- ch:
                if !ok {
                    fmt.Printf("Go channel closed - stop!\n");
                    return ret
                }
                ret = append(ret, s)
                //fmt.Printf("Got response: '%s'\n", s)
            //case <- timer.C:
                //fmt.Printf("Timeout!\n")
            //    return ret
        }
    }
    return ret
}

<commit_msg>Added parsing of incoming string<commit_after>package onkyo

import "github.com/cnf/go-claw/tools"
import "github.com/cnf/go-claw/clog"
import "strconv"
import "net"
import "time"
import "strings"

const ONKYO_PORT = 60128
const ONKYO_MAGIC = "!xECNQSTN"

type TargetDevice struct {
    Name string
    Model string
    Params map[string]string   // Parameters that should be saved in the config file
    Detected map[string]string // Detected parameters that should be re-evaluated each connect
}

func runOnkyoDetect(ch chan *TargetDevice, timeout int) {
    defer close(ch)
    addrs := tools.BroadcastAddrs()
    port := strconv.Itoa(ONKYO_PORT)

    cmd := &OnkyoCommandTCP{ONKYO_MAGIC}
    transmit_str, err := cmd.Bytes()
    if err != nil {
        return
    }

    udpaddr, err := net.ResolveUDPAddr("udp4", ":0")
    if err != nil {
        return
    }
    conn, err := net.ListenUDP("udp4", udpaddr)
    if err != nil {
        clog.Error("targets/onkyo: Could not listen on UDP Address %v: %s", udpaddr, err.Error())
        return
    }
    defer conn.Close()
    conn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Millisecond))

    // Send broadcast
    for _, a := range addrs {
        udpdest, err := net.ResolveUDPAddr("udp4", a+":"+port)
        if err != nil {
            clog.Warn("targets/onkyo: ResolveUDPAddr('udp4','%s:%s'): %s", a, port, err.Error())
            continue
        }
        conn.WriteToUDP(transmit_str, udpdest)
        // Ignore errors - just continue
    }

    // Now loop listening to responses with timeout
    buf := make([]byte, 1024)
    for {
        rlen, raddr, err := conn.ReadFromUDP(buf);
        if err != nil {
            if nerr, ok := err.(net.Error); !ok || !nerr.Temporary() || nerr.Timeout() {
                if (!nerr.Timeout()) {
                    // Unexpected!
                    clog.Error("targets/onkyo: ReadFromUDP: %s", err.Error())
                }
                return
            }
            continue
        }
        //fmt.Printf("Got response from: %v (len: %d):\n", raddr, rlen)
        resp := &OnkyoCommandTCP{}
        if err := resp.Parse(buf[0:rlen]); err != nil {
            // fmt.Printf("Error parsing Onkyo response: %s", err.Error())
            clog.Warn("targets/onkyo: Parse Onkyo Response: %s", err.Error())
            continue
        }
        // Parse the message:
        // ECN<model name>/<ISCP port>/<region:DX|XX|JJ>/<id>
        rmsg := resp.Message()
        if len(rmsg) <= 3 {
            clog.Warn("targets/onkyo: Autodetect response too short: %d", len(rmsg))
            continue
        }
        if rmsg[0:3] != "ECN" {
            clog.Warn("targets/onkyo: Autodetect Unexpected response: %s", rmsg)
            continue
        }
        splitmsg := strings.Split(rmsg[3:], "/")
        if len(splitmsg) != 4 {
            clog.Warn("targets/onkyo: Autodetect invalid response format, expected 4 fields: %s", rmsg)
            continue
        }
        tgt := &TargetDevice{}
        tgt.Name = splitmsg[0] + " (" + splitmsg[3] + ")"
        tgt.Model = splitmsg[0]
        tgt.Params = make(map[string]string)
        tgt.Params["id"] = splitmsg[3]
        tgt.Params["Model"] = splitmsg[0]
        tgt.Detected = make(map[string]string)
        tgt.Detected["port"] = splitmsg[1]
        host, _, err := net.SplitHostPort(raddr.String())
        if err != nil {
            clog.Warn("Could not split the host/port for '%s'", raddr.String())
            continue
        }
        tgt.Detected["ip"] = host

        ch <- tgt

        // Once we got a response, wait maximum 50ms for other responses
        conn.SetReadDeadline(time.Now().Add(time.Duration(20) * time.Millisecond))
    }
}

func OnkyoAutoDetect(timeout int) []TargetDevice {
    ch := make(chan *TargetDevice)
    go runOnkyoDetect(ch, timeout)
    ret := make([]TargetDevice, 0)
    //timer := time.NewTimer(time.Millisecond * time.Duration(timeout + 200))
    for {
        select {
            case s, ok := <- ch:
                if !ok {
                    return ret
                }
                ret = append(ret, *s)
            //case <- timer.C:
                //fmt.Printf("Timeout!\n")
                //return ret
        }
    }
    return ret
}

<|endoftext|>


<commit_before>package file

import (
	"container/list"
	"sync"
	"time"

	"github.com/wx13/sith/file/buffer"
	"github.com/wx13/sith/file/cursor"
)

type BufferState struct {
	buff  buffer.Buffer
	mc    cursor.MultiCursor
	saved bool
}

func NewBufferState(buff buffer.Buffer, mc cursor.MultiCursor) *BufferState {
	return &BufferState{
		buff: buff,
		mc:   mc,
	}
}

type BufferHist struct {
	list      *list.List
	element   *list.Element
	elemMutex *sync.Mutex

	snapChan chan struct{}
	snapReq  SnapshotRequest
	reqMutex *sync.Mutex
}

type SnapshotRequest struct {
	Buffer buffer.Buffer
	Cursor cursor.MultiCursor
}

func NewBufferHist(buffer buffer.Buffer, cursor cursor.MultiCursor) *BufferHist {
	bh := BufferHist{}
	state := NewBufferState(buffer, cursor)
	bh.list = list.New()
	bh.element = bh.list.PushBack(state)
	bh.elemMutex = &sync.Mutex{}
	bh.snapChan = make(chan struct{}, 1)
	bh.reqMutex = &sync.Mutex{}
	bh.handleSnapshots()
	return &bh
}

func (bh *BufferHist) ForceSnapshot(buff buffer.Buffer, mc cursor.MultiCursor) {
	bh.snapshot(buff.Dup(), mc.Dup())
}

func (bh *BufferHist) Snapshot(buff buffer.Buffer, mc cursor.MultiCursor) {
	request := SnapshotRequest{
		Buffer: buff.Dup(),
		Cursor: mc.Dup(),
	}

	bh.reqMutex.Lock()
	bh.snapReq = request
	bh.reqMutex.Unlock()

	select {
	case bh.snapChan <- struct{}{}:
	default:
	}
}

func (bh *BufferHist) handleSnapshots() {
	go func() {
		for range time.Tick(time.Millisecond * 100) {
			select {
			case <-bh.snapChan:
				bh.reqMutex.Lock()
				bh.snapshot(bh.snapReq.Buffer, bh.snapReq.Cursor)
				bh.reqMutex.Unlock()
			}
		}
	}()
}

func (bh *BufferHist) SnapshotSaved() {
	bh.elemMutex.Lock()
	bh.element.Value.(*BufferState).saved = true
	bh.elemMutex.Unlock()
}

func (bh *BufferHist) snapshot(buff buffer.Buffer, mc cursor.MultiCursor) {

	state := NewBufferState(buff, mc)
	bh.elemMutex.Lock()
	bh.element = bh.list.InsertAfter(state, bh.element)
	bh.elemMutex.Unlock()
	bh.Trim()
}

func (bh *BufferHist) Trim() {

	if bh.list.Len() < 200 {
		return
	}

	rm := []*list.Element{}

	n := 0
	ns := 0
	for el := bh.element.Next(); el != nil; el = el.Next() {
		n++
		state := el.Value.(*BufferState)
		if state.saved {
			ns++
			if ns > 50 {
				rm = append(rm, el)
			}
			continue
		}
		if n > 50 {
			rm = append(rm, el)
		}
	}

	n = 0
	ns = 0
	for el := bh.element.Prev(); el != nil; el = el.Prev() {
		n++
		state := el.Value.(*BufferState)
		if state.saved {
			ns++
			if ns > 50 {
				rm = append(rm, el)
			}
			continue
		}
		if n > 50 {
			rm = append(rm, el)
		}
	}

	for _, el := range rm {
		bh.list.Remove(el)
	}

}

func (bh *BufferHist) Current() (buffer.Buffer, cursor.MultiCursor) {
	state := bh.element.Value.(*BufferState)
	return state.buff, state.mc.Dup()
}

func (bh *BufferHist) Next() (buffer.Buffer, cursor.MultiCursor) {
	next := bh.element.Next()
	if next != nil {
		bh.element = next
	}
	return bh.Current()
}

func (bh *BufferHist) Prev() (buffer.Buffer, cursor.MultiCursor) {
	bh.elemMutex.Lock()
	prev := bh.element.Prev()
	bh.elemMutex.Unlock()
	if prev != nil {
		bh.elemMutex.Lock()
		bh.element = prev
		bh.elemMutex.Unlock()
	}
	return bh.Current()
}

func (bh *BufferHist) NextSaved() (buffer.Buffer, cursor.MultiCursor) {
	for el := bh.element.Next(); el != nil; el = el.Next() {
		if el.Value.(*BufferState).saved {
			bh.SnapshotSaved()
			bh.elemMutex.Lock()
			bh.element = el
			bh.elemMutex.Unlock()
			break
		}
	}
	return bh.Current()
}

func (bh *BufferHist) PrevSaved() (buffer.Buffer, cursor.MultiCursor) {
	for el := bh.element.Prev(); el != nil; el = el.Prev() {
		if el.Value.(*BufferState).saved {
			bh.SnapshotSaved()
			bh.elemMutex.Lock()
			bh.element = el
			bh.elemMutex.Unlock()
			break
		}
	}
	return bh.Current()
}
<commit_msg>keep cursor on undo/redo change spot<commit_after>package file

import (
	"container/list"
	"sync"
	"time"

	"github.com/wx13/sith/file/buffer"
	"github.com/wx13/sith/file/cursor"
)

type BufferState struct {
	buff  buffer.Buffer
	mc    cursor.MultiCursor
	saved bool
}

func NewBufferState(buff buffer.Buffer, mc cursor.MultiCursor) *BufferState {
	return &BufferState{
		buff: buff,
		mc:   mc,
	}
}

type BufferHist struct {
	list      *list.List
	element   *list.Element
	elemMutex *sync.Mutex

	snapChan chan struct{}
	snapReq  SnapshotRequest
	reqMutex *sync.Mutex
}

type SnapshotRequest struct {
	Buffer buffer.Buffer
	Cursor cursor.MultiCursor
}

func NewBufferHist(buffer buffer.Buffer, cursor cursor.MultiCursor) *BufferHist {
	bh := BufferHist{}
	state := NewBufferState(buffer, cursor)
	bh.list = list.New()
	bh.element = bh.list.PushBack(state)
	bh.elemMutex = &sync.Mutex{}
	bh.snapChan = make(chan struct{}, 1)
	bh.reqMutex = &sync.Mutex{}
	bh.handleSnapshots()
	return &bh
}

func (bh *BufferHist) ForceSnapshot(buff buffer.Buffer, mc cursor.MultiCursor) {
	bh.snapshot(buff.Dup(), mc.Dup())
}

func (bh *BufferHist) Snapshot(buff buffer.Buffer, mc cursor.MultiCursor) {
	request := SnapshotRequest{
		Buffer: buff.Dup(),
		Cursor: mc.Dup(),
	}

	bh.reqMutex.Lock()
	bh.snapReq = request
	bh.reqMutex.Unlock()

	select {
	case bh.snapChan <- struct{}{}:
	default:
	}
}

func (bh *BufferHist) handleSnapshots() {
	go func() {
		for range time.Tick(time.Millisecond * 100) {
			select {
			case <-bh.snapChan:
				bh.reqMutex.Lock()
				bh.snapshot(bh.snapReq.Buffer, bh.snapReq.Cursor)
				bh.reqMutex.Unlock()
			}
		}
	}()
}

func (bh *BufferHist) SnapshotSaved() {
	bh.elemMutex.Lock()
	bh.element.Value.(*BufferState).saved = true
	bh.elemMutex.Unlock()
}

func (bh *BufferHist) snapshot(buff buffer.Buffer, mc cursor.MultiCursor) {

	curBuf, curMC := bh.Current()
	curRow := curMC.GetRow(0)
	newRow := mc.GetRow(0)
	if curRow-newRow > 5 || newRow-curRow > 5 {
		state := NewBufferState(curBuf, mc)
		bh.elemMutex.Lock()
		bh.element = bh.list.InsertAfter(state, bh.element)
		bh.elemMutex.Unlock()
	}

	state := NewBufferState(buff, mc)
	bh.elemMutex.Lock()
	bh.element = bh.list.InsertAfter(state, bh.element)
	bh.elemMutex.Unlock()

	bh.Trim()

}

func (bh *BufferHist) Trim() {

	if bh.list.Len() < 200 {
		return
	}

	rm := []*list.Element{}

	n := 0
	ns := 0
	for el := bh.element.Next(); el != nil; el = el.Next() {
		n++
		state := el.Value.(*BufferState)
		if state.saved {
			ns++
			if ns > 50 {
				rm = append(rm, el)
			}
			continue
		}
		if n > 50 {
			rm = append(rm, el)
		}
	}

	n = 0
	ns = 0
	for el := bh.element.Prev(); el != nil; el = el.Prev() {
		n++
		state := el.Value.(*BufferState)
		if state.saved {
			ns++
			if ns > 50 {
				rm = append(rm, el)
			}
			continue
		}
		if n > 50 {
			rm = append(rm, el)
		}
	}

	for _, el := range rm {
		bh.list.Remove(el)
	}

}

func (bh *BufferHist) Current() (buffer.Buffer, cursor.MultiCursor) {
	state := bh.element.Value.(*BufferState)
	return state.buff, state.mc.Dup()
}

func (bh *BufferHist) Next() (buffer.Buffer, cursor.MultiCursor) {
	next := bh.element.Next()
	if next != nil {
		bh.element = next
	}
	return bh.Current()
}

func (bh *BufferHist) Prev() (buffer.Buffer, cursor.MultiCursor) {
	bh.elemMutex.Lock()
	prev := bh.element.Prev()
	bh.elemMutex.Unlock()
	if prev != nil {
		bh.elemMutex.Lock()
		bh.element = prev
		bh.elemMutex.Unlock()
	}
	return bh.Current()
}

func (bh *BufferHist) NextSaved() (buffer.Buffer, cursor.MultiCursor) {
	for el := bh.element.Next(); el != nil; el = el.Next() {
		if el.Value.(*BufferState).saved {
			bh.SnapshotSaved()
			bh.elemMutex.Lock()
			bh.element = el
			bh.elemMutex.Unlock()
			break
		}
	}
	return bh.Current()
}

func (bh *BufferHist) PrevSaved() (buffer.Buffer, cursor.MultiCursor) {
	for el := bh.element.Prev(); el != nil; el = el.Prev() {
		if el.Value.(*BufferState).saved {
			bh.SnapshotSaved()
			bh.elemMutex.Lock()
			bh.element = el
			bh.elemMutex.Unlock()
			break
		}
	}
	return bh.Current()
}
<|endoftext|>


<commit_before>// Copyright 2014 Unknwon
//
// 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 admin

import (
	"fmt"
	"os"
	"path"
	"time"

	"github.com/Unknwon/com"

	"github.com/gpmgo/switch/models"
	"github.com/gpmgo/switch/modules/archive"
	"github.com/gpmgo/switch/modules/middleware"
	"github.com/gpmgo/switch/modules/qiniu"
	"github.com/gpmgo/switch/modules/setting"
)

type ApiRevesion struct {
	Id       int64       `json:"id"`
	Package  *ApiPackage `json:"package"`
	Revision string      `json:"revision"`
	Size     int64       `json:"size"`
	Updated  time.Time   `json:"updated"`
}

type ApiPackage struct {
	Id         int64     `json:"id"`
	ImportPath string    `json:"import_path"`
	Created    time.Time `json:"created"`
}

func ListLargeRevisions(ctx *middleware.Context) {
	revs, err := models.GetLocalRevisions()
	if err != nil {
		ctx.JSON(500, map[string]string{
			"error": fmt.Sprintf("fail to get local revisions: %v", err),
		})
		return
	}

	largeRevs := make([]*ApiRevesion, 0, len(revs)/2)
	for _, rev := range revs {
		pkg, err := models.GetPakcageById(rev.PkgId)
		if err != nil {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to get package by ID(%d): %v", rev.PkgId, err),
			})
			return
		}

		ext := archive.GetExtension(pkg.ImportPath)
		localPath := path.Join(pkg.ImportPath, rev.Revision)
		fpath := path.Join(setting.ArchivePath, localPath+ext)

		if !com.IsFile(fpath) {
			continue
		}

		// Check archive size.
		f, err := os.Open(fpath)
		if err != nil {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to open file(%s): %v", fpath, err),
			})
			return
		}
		fi, err := f.Stat()
		if err != nil {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to get file info(%s): %v", fpath, err),
			})
			return
		}
		// Greater then MAX_UPLOAD_SIZE.
		if fi.Size() > setting.MaxUploadSize<<20 {
			largeRevs = append(largeRevs, &ApiRevesion{
				Id: rev.Id,
				Package: &ApiPackage{
					Id:         pkg.Id,
					ImportPath: pkg.ImportPath,
					Created:    pkg.Created,
				},
				Revision: rev.Revision,
				Size:     fi.Size(),
				Updated:  rev.Updated,
			})
			continue
		}
	}

	ctx.JSON(200, map[string]interface{}{
		"revisions": &largeRevs,
	})
}

func BlockPackage(ctx *middleware.Context) {
	id := ctx.QueryInt64("id")
	pkg, err := models.GetPakcageById(id)
	if err != nil {
		if err == models.ErrPackageNotExist {
			ctx.JSON(404, map[string]string{
				"error": err.Error(),
			})
		} else {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to get package by ID(%d): %v", id, err),
			})
		}
		return
	}

	revs, err := pkg.GetRevisions()
	if err != nil {
		ctx.JSON(500, map[string]string{
			"error": fmt.Sprintf("fail to get package revisions by ID(%d): %v", id, err),
		})
		return
	}

	models.BlockPackage(pkg, revs, ctx.Query("note"))
	return

	// Delete package archives.
	ext := archive.GetExtension(pkg.ImportPath)
	for _, rev := range revs {
		switch rev.Storage {
		case models.LOCAL:
			localPath := path.Join(pkg.ImportPath, rev.Revision)
			os.Remove(path.Join(setting.ArchivePath, localPath+ext))
		case models.QINIU:
			key := pkg.ImportPath + "-" + rev.Revision + ext
			if err = qiniu.DeleteArchive(key); err != nil {
				ctx.JSON(500, map[string]string{
					"error": fmt.Sprintf("fail to delete archive(%s): %v", key, err),
				})
				return
			}
		}
	}
}
<commit_msg>fix bug<commit_after>// Copyright 2014 Unknwon
//
// 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 admin

import (
	"fmt"
	"os"
	"path"
	"time"

	"github.com/Unknwon/com"

	"github.com/gpmgo/switch/models"
	"github.com/gpmgo/switch/modules/archive"
	"github.com/gpmgo/switch/modules/middleware"
	"github.com/gpmgo/switch/modules/qiniu"
	"github.com/gpmgo/switch/modules/setting"
)

type ApiRevesion struct {
	Id       int64       `json:"id"`
	Package  *ApiPackage `json:"package"`
	Revision string      `json:"revision"`
	Size     int64       `json:"size"`
	Updated  time.Time   `json:"updated"`
}

type ApiPackage struct {
	Id         int64     `json:"id"`
	ImportPath string    `json:"import_path"`
	Created    time.Time `json:"created"`
}

func ListLargeRevisions(ctx *middleware.Context) {
	revs, err := models.GetLocalRevisions()
	if err != nil {
		ctx.JSON(500, map[string]string{
			"error": fmt.Sprintf("fail to get local revisions: %v", err),
		})
		return
	}

	largeRevs := make([]*ApiRevesion, 0, len(revs)/2)
	for _, rev := range revs {
		pkg, err := models.GetPakcageById(rev.PkgId)
		if err != nil {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to get package by ID(%d): %v", rev.PkgId, err),
			})
			return
		}

		ext := archive.GetExtension(pkg.ImportPath)
		localPath := path.Join(pkg.ImportPath, rev.Revision)
		fpath := path.Join(setting.ArchivePath, localPath+ext)

		if !com.IsFile(fpath) {
			continue
		}

		// Check archive size.
		f, err := os.Open(fpath)
		if err != nil {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to open file(%s): %v", fpath, err),
			})
			return
		}
		fi, err := f.Stat()
		if err != nil {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to get file info(%s): %v", fpath, err),
			})
			return
		}
		// Greater then MAX_UPLOAD_SIZE.
		if fi.Size() > setting.MaxUploadSize<<20 {
			largeRevs = append(largeRevs, &ApiRevesion{
				Id: rev.Id,
				Package: &ApiPackage{
					Id:         pkg.Id,
					ImportPath: pkg.ImportPath,
					Created:    pkg.Created,
				},
				Revision: rev.Revision,
				Size:     fi.Size(),
				Updated:  rev.Updated,
			})
			continue
		}
	}

	ctx.JSON(200, map[string]interface{}{
		"revisions": &largeRevs,
	})
}

func BlockPackage(ctx *middleware.Context) {
	id := ctx.QueryInt64("id")
	pkg, err := models.GetPakcageById(id)
	if err != nil {
		if err == models.ErrPackageNotExist {
			ctx.JSON(404, map[string]string{
				"error": err.Error(),
			})
		} else {
			ctx.JSON(500, map[string]string{
				"error": fmt.Sprintf("fail to get package by ID(%d): %v", id, err),
			})
		}
		return
	}

	revs, err := pkg.GetRevisions()
	if err != nil {
		ctx.JSON(500, map[string]string{
			"error": fmt.Sprintf("fail to get package revisions by ID(%d): %v", id, err),
		})
		return
	}

	// Delete package archives.
	ext := archive.GetExtension(pkg.ImportPath)
	for _, rev := range revs {
		switch rev.Storage {
		case models.QINIU:
			key := pkg.ImportPath + "-" + rev.Revision + ext
			if err = qiniu.DeleteArchive(key); err != nil {
				ctx.JSON(500, map[string]string{
					"error": fmt.Sprintf("fail to delete archive(%s): %v", key, err),
				})
				return
			}
		}
	}
	os.RemoveAll(path.Join(setting.ArchivePath, pkg.ImportPath))

	if err = models.BlockPackage(pkg, revs, ctx.Query("note")); err != nil {
		ctx.JSON(500, map[string]string{
			"error": fmt.Sprintf("fail to block package by ID(%d): %v", id, err),
		})
		return
	}

	ctx.JSON(200, map[string]interface{}{
		"ok": true,
	})
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"

	"github.com/gorilla/mux"
)

type LinkLedger struct {
	Source int `json:"source"` // the source id
	Block  int `json:"block"`  // the block id
	Id     int `json:"id"`     // link id
}

type ProtoLink struct {
	Source int `json:"source"` // the source id
	Block  int `json:"block"`  // the block id
}

func (s *Server) CreateLink(l ProtoLink) (*LinkLedger, error) {
	b, ok := s.blocks[l.Block]
	if !ok {
		return nil, errors.New("could not find block")
	}

	sl, ok := s.sources[l.Source]
	if !ok {
		return nil, errors.New("could not find source")
	}

	link := &LinkLedger{
		Source: l.Source,
		Block:  l.Block,
		Id:     s.GetNextID(),
	}

	if b.GetParent() != sl.GetParent() {
		return nil, errors.New("block and source must be in the same group, cannot link")
	}

	err := b.Block.SetSource(sl.Source)
	if err != nil {
		return nil, err
	}

	s.links[link.Id] = link

	s.websocketBroadcast(Update{Action: CREATE, Type: LINK, Data: link})

	return link, nil
}

func (s *Server) DeleteLink(id int) error {
	link, ok := s.links[id]
	if !ok {
		return errors.New("could not find link")
	}

	s.blocks[link.Block].Block.SetSource(nil)
	delete(s.links, id)

	s.websocketBroadcast(Update{Action: CREATE, Type: CONNECTION, Data: link})
	return nil
}

func (s *Server) ListLinks() []LinkLedger {
	links := []LinkLedger{}
	for _, l := range s.links {
		links = append(links, *l)
	}
	return links
}

func (s *Server) LinkIndexHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	c := s.ListLinks()

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(c); err != nil {
		panic(err)
	}
}

func (s *Server) LinkCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	var newLink ProtoLink
	json.Unmarshal(body, &newLink)

	s.Lock()
	defer s.Unlock()

	nl, err := s.CreateLink(newLink)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, nl)
}

func (s *Server) LinkDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteLink(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
<commit_msg>fixing link api<commit_after>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"

	"github.com/gorilla/mux"
)

type LinkLedger struct {
	Source struct {
		Id int `json:"id"`
	} `json:"source"` // the soure id
	Block struct {
		Id int `json:"id"`
	} `json:"block"` // the block id
	Id int `json:"id"` // link id
}

type ProtoLink struct {
	Source struct {
		Id int `json:"id"`
	} `json:"source"` // the soure id
	Block struct {
		Id int `json:"id"`
	} `json:"block"` // the block id
}

func (s *Server) CreateLink(l ProtoLink) (*LinkLedger, error) {
	b, ok := s.blocks[l.Block.Id]
	if !ok {
		return nil, errors.New("could not find block")
	}

	sl, ok := s.sources[l.Source.Id]
	if !ok {
		return nil, errors.New("could not find source")
	}

	link := &LinkLedger{}
	link.Id = s.GetNextID()
	link.Source.Id = l.Source.Id
	link.Block.Id = l.Block.Id

	if b.GetParent() != sl.GetParent() {
		return nil, errors.New("block and source must be in the same group, cannot link")
	}

	err := b.Block.SetSource(sl.Source)
	if err != nil {
		return nil, err
	}

	s.links[link.Id] = link

	s.websocketBroadcast(Update{Action: CREATE, Type: LINK, Data: wsLink{*link}})

	return link, nil
}

func (s *Server) DeleteLink(id int) error {
	link, ok := s.links[id]
	if !ok {
		return errors.New("could not find link")
	}

	s.blocks[link.Block.Id].Block.SetSource(nil)
	delete(s.links, id)

	s.websocketBroadcast(Update{Action: DELETE, Type: LINK, Data: wsLink{wsId{id}}})
	return nil
}

func (s *Server) ListLinks() []LinkLedger {
	links := []LinkLedger{}
	for _, l := range s.links {
		links = append(links, *l)
	}
	return links
}

func (s *Server) LinkIndexHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	c := s.ListLinks()

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(c); err != nil {
		panic(err)
	}
}

func (s *Server) LinkCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	var newLink ProtoLink
	json.Unmarshal(body, &newLink)

	s.Lock()
	defer s.Unlock()

	nl, err := s.CreateLink(newLink)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, nl)
}

func (s *Server) LinkDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteLink(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/TopHatCroat/CryptoChat-server/constants"
	"github.com/TopHatCroat/CryptoChat-server/database"
	"github.com/TopHatCroat/CryptoChat-server/models"
	"github.com/TopHatCroat/CryptoChat-server/protocol"
	"github.com/TopHatCroat/CryptoChat-server/tools"
	"github.com/gorilla/context"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

type AppHandler func(http.ResponseWriter, *http.Request) *appError

type Adapter func(http.Handler) http.Handler

type appError struct {
	Error   error
	Message string
	Code    int
}

func Logger(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		log := models.Log{
			SourceAddr: req.RemoteAddr,
			Params:     req.RequestURI,
			Method:     req.Method,
			Cipher:     req.TLS.CipherSuite,
			Timestamp:  time.Now().UnixNano(),
		}

		if err := log.Log(); err != nil {
			Respond(rw, map[string]string{"error": constants.REQUEST_REJECTED})
			return
		}

		handler.ServeHTTP(rw, req)

		log.TimeRequest()

		if err := GetError(req); err != nil {
			Respond(rw, map[string]string{"error": constants.REQUEST_REJECTED})
			return
		}
	})
}

func JSONDecoder(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		decoder := json.NewDecoder(req.Body)
		//var msg json.RawMessage
		fullMsg := protocol.CompleteMessage{
		//Content: &msg,
		}
		if err := decoder.Decode(&fullMsg); err != nil {
			context.Set(req, "err", err)
			return
		}

		context.Set(req, "json", fullMsg)

		handler.ServeHTTP(rw, req)
	})
}

func (fn AppHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if err := fn(rw, req); err != nil {
		Respond(rw, map[string]string{"error": err.Message})
	}
}

func GetJSON(r *http.Request) protocol.CompleteMessage {
	if rv := context.Get(r, "json"); rv != nil {
		return rv.(protocol.CompleteMessage)
	}
	return protocol.CompleteMessage{}
}

func GetError(r *http.Request) error {
	if rv := context.Get(r, "err"); rv != nil {
		return rv.(error)
	}
	return nil
}

func Respond(rw http.ResponseWriter, data interface{}) {
	encoder := json.NewEncoder(rw)
	encoder.Encode(data)
}

func sendHandler(rw http.ResponseWriter, req *http.Request) *appError {
	fullMsg := GetJSON(req)

	user, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return &appError{err, err.Error(), 500}
	}

	if fullMsg.Type == "S" {
		var messageRequest protocol.Message
		log.Println("Recieved message request")
		json.Unmarshal(*fullMsg.Content, &messageRequest)

		message := &models.Message{
			RecieverID: int64(messageRequest.Receiver),
			SenderID:   user.ID,
			Content:    messageRequest.Content,
			KeyHash:    messageRequest.KeyHash,
			CreatedAt:  messageRequest.Timestamp,
		}

		err := message.Save()
		if err != nil {
			return &appError{err, err.Error(), 500}
		}

		var messageResponse protocol.MessageResponse
		messageResponse.Message = constants.MESSAGE_SENT
		Respond(rw, messageResponse)
	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return &appError{err, err.Error(), 500}
	}

	return nil
}

func messagesHandler(rw http.ResponseWriter, req *http.Request) *appError {
	fullMsg := GetJSON(req)

	user, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return &appError{err, err.Error(), 500}
	}

	if fullMsg.Type == "M" {
		var getMessagesRequest protocol.GetMessagesRequest
		json.Unmarshal(*fullMsg.Content, &getMessagesRequest)

		messages, err := models.GetNewMessagesForUser(user, getMessagesRequest.LastMessageTimestamp)
		if err != nil {
			return &appError{err, err.Error(), 500}
		}

		var getMessagesResponse protocol.GetMessagesResponse
		getMessagesResponse.Messages = messages

		Respond(rw, getMessagesResponse)
	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return &appError{err, err.Error(), 500}
	}

	return nil
}

func keyHandler(rw http.ResponseWriter, req *http.Request) *appError {
	fullMsg := GetJSON(req)

	_, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return &appError{err, err.Error(), 401}
	}

	//key submit
	if fullMsg.Type == "K" {
		var keyRequest protocol.KeyRequest
		json.Unmarshal(*fullMsg.Content, &keyRequest)

		key := &models.Key{
			Key:       keyRequest.Key,
			Hash:      keyRequest.Hash,
			FriendID:  keyRequest.UserID,
			CreatedAt: keyRequest.CreatedAt,
		}

		err := key.Save()
		if err != nil {
			return &appError{err, err.Error(), 500}
		}

		var keyResponse protocol.KeyResponse
		keyResponse.Status = constants.KEY_SUBMIT_SUCCESS
		keyResponse.Hash = key.Hash

		Respond(rw, keyResponse)
	} else if fullMsg.Type == "KR" {
		var keyRequest protocol.KeyRequest
		json.Unmarshal(*fullMsg.Content, &keyRequest)

		key, err := models.FindKeyByHash(keyRequest.Hash)
		if err != nil {
			return &appError{err, err.Error(), 400}
		}

		var keyResponse protocol.KeyResponse
		keyResponse.Status = constants.KEY_FOUND_SUCCESS
		keyResponse.Key = key.Key
		keyResponse.Hash = key.Hash
		keyResponse.UserID = key.FriendID
		keyResponse.CreatedAt = key.CreatedAt

		Respond(rw, keyResponse)
	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return &appError{err, err.Error(), 500}
	}

	return nil
}

func userHandler(rw http.ResponseWriter, req *http.Request) *appError {
	fullMsg := GetJSON(req)

	_, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return &appError{err, err.Error(), 500}
	}

	if fullMsg.Type == "U" {
		var friendRequest protocol.FriendRequest
		json.Unmarshal(*fullMsg.Content, &friendRequest)

		user, err := models.FindUserByCreds(friendRequest.Username)
		if err != nil {
			return &appError{err, err.Error(), 500}
		}

		var friendResponse protocol.FriendResponse
		friendResponse.User.APIID = user.ID
		friendResponse.User.Username = user.Username
		friendResponse.User.PublicKey = user.PublicKey

		Respond(rw, friendResponse)
	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return &appError{err, err.Error(), 500}
	}

	return nil
}

func registerHandler(rw http.ResponseWriter, req *http.Request) *appError {
	fullMsg := GetJSON(req)
	if fullMsg.Type == "R" {
		var connectRequest protocol.ConnectRequest
		log.Println("Recieved register request")
		json.Unmarshal(*fullMsg.Content, &connectRequest)
		var user models.User
		user, err := models.CreateUser(connectRequest.UserName, connectRequest.Password, connectRequest.PublicKey)
		if err != nil {
			return &appError{err, err.Error(), 500}
		}

		var connectResponse protocol.ConnectResponse
		connectResponse.Type = constants.REGISTER_SUCCESS
		connectResponse.Token = user.Username
		Respond(rw, connectResponse)

	} else if fullMsg.Type == "L" {
		var connectRequest protocol.ConnectRequest
		log.Println("Recieved login request")
		json.Unmarshal(*fullMsg.Content, &connectRequest)
		var user models.User
		user, err := models.FindUserByCreds(connectRequest.UserName)
		if err != nil {
			return &appError{err, err.Error(), 500}
		}

		userToken, err := user.LogIn(connectRequest.Password)
		if err != nil {
			return &appError{err, err.Error(), 500}
		}

		var connectResponse protocol.ConnectResponse
		connectResponse.Type = constants.LOGIN_SUCCESS
		connectResponse.Token = userToken
		Respond(rw, connectResponse)

	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return &appError{err, err.Error(), 500}
	}

	return nil
}

func init() {
	if _, err := os.Stat(constants.TOKEN_KEY_FILE); os.IsNotExist(err) {
		tools.GenerateTokenKey()
	}

	if err := os.Setenv(constants.EDITION_VAR, constants.SERVER_EDITION); err != nil {
		panic(err)
	}
}

func main() {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigs
		fmt.Println()
		fmt.Println(sig)
		database.CloseDatabase()
		fmt.Println("SecureChat server closed...")
		os.Exit(0)
	}()

	configuration := &tls.Config{
		MinVersion:               tls.VersionTLS12,
		CurvePreferences:         []tls.CurveID{tls.CurveP521},
		PreferServerCipherSuites: true,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
			tls.TLS_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_RSA_WITH_AES_256_CBC_SHA,
		},
	}

	mux := http.NewServeMux()
	mux.Handle("/register", Logger(JSONDecoder(AppHandler(registerHandler))))
	mux.Handle("/send", Logger(JSONDecoder(AppHandler(sendHandler))))
	mux.Handle("/messages", Logger(JSONDecoder(AppHandler(messagesHandler))))
	mux.Handle("/user", Logger(JSONDecoder(AppHandler(userHandler))))
	mux.Handle("/keys", Logger(JSONDecoder(AppHandler(keyHandler))))

	server := &http.Server{
		Addr:         ":44333",
		Handler:      mux,
		TLSConfig:    configuration,
		TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler), 0),
	}

	database.GetDatabase()
	log.Fatal(server.ListenAndServeTLS("server.cert", "server.key"))

}
<commit_msg>simplified server responding<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/TopHatCroat/CryptoChat-server/constants"
	"github.com/TopHatCroat/CryptoChat-server/database"
	"github.com/TopHatCroat/CryptoChat-server/models"
	"github.com/TopHatCroat/CryptoChat-server/protocol"
	"github.com/TopHatCroat/CryptoChat-server/tools"
	"github.com/gorilla/context"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

type AppHandler func(http.ResponseWriter, *http.Request) *appData

type Adapter func(http.Handler) http.Handler

type appData struct {
	Data   interface{}
	Error  error
	Status string
	Code   int
}

func Logger(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		log := models.Log{
			SourceAddr: req.RemoteAddr,
			Params:     req.RequestURI,
			Method:     req.Method,
			Cipher:     req.TLS.CipherSuite,
			Timestamp:  time.Now().UnixNano(),
		}

		if err := log.Log(); err != nil {
			encoder := json.NewEncoder(rw)
			encoder.Encode(map[string]string{"error": constants.REQUEST_REJECTED})
			return
		}

		handler.ServeHTTP(rw, req)

		log.TimeRequest()

		if err := GetError(req); err != nil {
			encoder := json.NewEncoder(rw)
			encoder.Encode(map[string]string{"error": constants.REQUEST_REJECTED})
			return
		}
	})
}

func JSONDecoder(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		decoder := json.NewDecoder(req.Body)
		//var msg json.RawMessage
		fullMsg := protocol.CompleteMessage{
		//Content: &msg,
		}
		if err := decoder.Decode(&fullMsg); err != nil {
			context.Set(req, "err", err)
			return
		}

		context.Set(req, "json", fullMsg)

		handler.ServeHTTP(rw, req)
	})
}

func (fn AppHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	data := fn(rw, req)

	if data.Error != nil {
		encoder := json.NewEncoder(rw)
		encoder.Encode(map[string]string{"error": data.Error.Error(), "status": data.Status})
	} else {
		encoder := json.NewEncoder(rw)
		encoder.Encode(data.Data)
	}
}

func GetJSON(r *http.Request) protocol.CompleteMessage {
	if rv := context.Get(r, "json"); rv != nil {
		return rv.(protocol.CompleteMessage)
	}
	return protocol.CompleteMessage{}
}

func GetError(r *http.Request) error {
	if rv := context.Get(r, "err"); rv != nil {
		return rv.(error)
	}
	return nil
}

func validResponse(data interface{}) *appData {
	return &appData{
		Data: data,
	}
}

func errorResponse(err error, status string, code int) *appData {
	return &appData{
		Error:  err,
		Status: status,
		Code:   code,
	}
}

func sendHandler(rw http.ResponseWriter, req *http.Request) *appData {
	fullMsg := GetJSON(req)

	user, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return errorResponse(err, err.Error(), 500)
	}

	if fullMsg.Type == "S" {
		var messageRequest protocol.Message
		log.Println("Recieved message request")
		json.Unmarshal(*fullMsg.Content, &messageRequest)

		message := &models.Message{
			RecieverID: int64(messageRequest.Receiver),
			SenderID:   user.ID,
			Content:    messageRequest.Content,
			KeyHash:    messageRequest.KeyHash,
			CreatedAt:  messageRequest.Timestamp,
		}

		err := message.Save()
		if err != nil {
			return errorResponse(err, err.Error(), 500)
		}

		var messageResponse protocol.MessageResponse
		messageResponse.Message = constants.MESSAGE_SENT
		return validResponse(messageResponse)

	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return errorResponse(err, err.Error(), 500)
	}
}

func messagesHandler(rw http.ResponseWriter, req *http.Request) *appData {
	fullMsg := GetJSON(req)

	user, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return errorResponse(err, err.Error(), 500)
	}

	if fullMsg.Type == "M" {
		var getMessagesRequest protocol.GetMessagesRequest
		json.Unmarshal(*fullMsg.Content, &getMessagesRequest)

		messages, err := models.GetNewMessagesForUser(user, getMessagesRequest.LastMessageTimestamp)
		if err != nil {
			return errorResponse(err, err.Error(), 500)
		}

		var getMessagesResponse protocol.GetMessagesResponse
		getMessagesResponse.Messages = messages

		return validResponse(getMessagesResponse)
	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return errorResponse(err, err.Error(), 500)
	}
}

func keyHandler(rw http.ResponseWriter, req *http.Request) *appData {
	fullMsg := GetJSON(req)

	_, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return errorResponse(err, err.Error(), 500)
	}

	//key submit
	if fullMsg.Type == "K" {
		var keyRequest protocol.KeyRequest
		json.Unmarshal(*fullMsg.Content, &keyRequest)

		key := &models.Key{
			Key:       keyRequest.Key,
			Hash:      keyRequest.Hash,
			FriendID:  keyRequest.UserID,
			CreatedAt: keyRequest.CreatedAt,
		}

		err := key.Save()
		if err != nil {
			return errorResponse(err, err.Error(), 500)
		}

		var keyResponse protocol.KeyResponse
		keyResponse.Status = constants.KEY_SUBMIT_SUCCESS
		keyResponse.Hash = key.Hash

		return validResponse(keyResponse)
	} else if fullMsg.Type == "KR" {
		var keyRequest protocol.KeyRequest
		json.Unmarshal(*fullMsg.Content, &keyRequest)

		key, err := models.FindKeyByHash(keyRequest.Hash)
		if err != nil {
			return errorResponse(err, err.Error(), 400)
		}

		var keyResponse protocol.KeyResponse
		keyResponse.Status = constants.KEY_FOUND_SUCCESS
		keyResponse.Key = key.Key
		keyResponse.Hash = key.Hash
		keyResponse.UserID = key.FriendID
		keyResponse.CreatedAt = key.CreatedAt

		return validResponse(keyResponse)
	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return errorResponse(err, err.Error(), 500)
	}
}

func userHandler(rw http.ResponseWriter, req *http.Request) *appData {
	fullMsg := GetJSON(req)

	_, err := models.FindUserByToken(fullMsg.Meta.Token)
	if err != nil {
		return errorResponse(err, err.Error(), 500)
	}

	if fullMsg.Type == "U" {
		var friendRequest protocol.FriendRequest
		json.Unmarshal(*fullMsg.Content, &friendRequest)

		user, err := models.FindUserByCreds(friendRequest.Username)
		if err != nil {
			return errorResponse(err, err.Error(), 500)
		}

		var friendResponse protocol.FriendResponse
		friendResponse.User.APIID = user.ID
		friendResponse.User.Username = user.Username
		friendResponse.User.PublicKey = user.PublicKey

		return validResponse(friendResponse)
	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return errorResponse(err, err.Error(), 500)
	}

	return nil
}

func registerHandler(rw http.ResponseWriter, req *http.Request) *appData {
	fullMsg := GetJSON(req)
	if fullMsg.Type == "R" {
		var connectRequest protocol.ConnectRequest
		log.Println("Recieved register request")
		json.Unmarshal(*fullMsg.Content, &connectRequest)
		var user models.User
		user, err := models.CreateUser(connectRequest.UserName, connectRequest.Password, connectRequest.PublicKey)
		if err != nil {
			return errorResponse(err, err.Error(), 500)
		}

		var connectResponse protocol.ConnectResponse
		connectResponse.Type = constants.REGISTER_SUCCESS
		connectResponse.Token = user.Username
		return validResponse(connectResponse)

	} else if fullMsg.Type == "L" {
		var connectRequest protocol.ConnectRequest
		log.Println("Recieved login request")
		json.Unmarshal(*fullMsg.Content, &connectRequest)
		var user models.User
		user, err := models.FindUserByCreds(connectRequest.UserName)
		if err != nil {
			return errorResponse(err, err.Error(), 500)
		}

		userToken, err := user.LogIn(connectRequest.Password)
		if err != nil {
			return errorResponse(err, err.Error(), 500)
		}

		var connectResponse protocol.ConnectResponse
		connectResponse.Type = constants.LOGIN_SUCCESS
		connectResponse.Token = userToken
		return validResponse(connectResponse)

	} else {
		err := errors.New(constants.WRONG_REQUEST)
		return errorResponse(err, err.Error(), 500)
	}
}

func init() {
	if _, err := os.Stat(constants.TOKEN_KEY_FILE); os.IsNotExist(err) {
		tools.GenerateTokenKey()
	}

	if err := os.Setenv(constants.EDITION_VAR, constants.SERVER_EDITION); err != nil {
		panic(err)
	}
}

func main() {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigs
		fmt.Println()
		fmt.Println(sig)
		database.CloseDatabase()
		fmt.Println("SecureChat server closed...")
		os.Exit(0)
	}()

	configuration := &tls.Config{
		MinVersion:               tls.VersionTLS12,
		CurvePreferences:         []tls.CurveID{tls.CurveP521},
		PreferServerCipherSuites: true,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
			tls.TLS_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_RSA_WITH_AES_256_CBC_SHA,
		},
	}

	mux := http.NewServeMux()
	mux.Handle("/register", Logger(JSONDecoder(AppHandler(registerHandler))))
	mux.Handle("/send", Logger(JSONDecoder(AppHandler(sendHandler))))
	mux.Handle("/messages", Logger(JSONDecoder(AppHandler(messagesHandler))))
	mux.Handle("/user", Logger(JSONDecoder(AppHandler(userHandler))))
	mux.Handle("/keys", Logger(JSONDecoder(AppHandler(keyHandler))))

	server := &http.Server{
		Addr:         ":44333",
		Handler:      mux,
		TLSConfig:    configuration,
		TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler), 0),
	}

	database.GetDatabase()
	log.Fatal(server.ListenAndServeTLS("server.cert", "server.key"))

}
<|endoftext|>


<commit_before>package main

import (
    "os"
    "io"
    "net/http"
    "github.com/go-martini/martini"
    "github.com/martini-contrib/render"
//  "github.com/garyburd/redigo/redis"
)

////////////////
//External API//
////////////////

func RestAddPkg(ren render.Render,req *http.Request,env *APIEnv) {
    ret := map[string]interface{}{"error":""}
    defer ren.JSON(200,ret)

    pkglen := req.ContentLength
    if pkglen <= 0 {
	ret["error"] = "EINVAL"
	return
    }
    if pkglen > PKG_MAXSIZE {
	ret["error"] = "ELIMIT"
	return
    }

    pkgid := PackageGenID()
    pkgfpath := STORAGE_PATH + "/package/" + pkgid + ".tar.xz"
    pkgfile,err := os.Create(pkgfpath)
    if err != nil {
	ret["error"] = "EIO"
	return
    }

    retlen,err := io.CopyN(pkgfile,req.Body,pkglen)
    pkgfile.Close()
    if retlen != pkglen || err != nil {
	os.Remove(pkgfpath)
	ret["error"] = "EIO"
	return
    }

    pkg := PackageCreate(pkgid,env)
    if pkg.Import(pkgfpath) != nil {
	os.Remove(pkgfpath)
	ret["error"] = "EINVAL"
	return
    }

    ret["pkgid"] = pkg.pkgid
}
func RestGetPkg(
    res http.ResponseWriter,
    req *http.Request,
    ram martini.Params,
    env *APIEnv,
) {
    pkg := PackageCreate(ram["pkgid"],env)
    err := pkg.Get()
    if err == nil && pkg.apiid == env.apiid {
	res.Header().Set("X-Accel-Redirect","/internal" + pkg.Export())
	res.WriteHeader(307)
	return
    } else if _,ok := err.(ErrPackageMiss); ok {
	if pkg.Transport() == nil {
	    res.Header().Set("X-Accel-Redirect","/internal" + pkg.Export())
	    res.WriteHeader(307)
	    return
	}
    }

    res.WriteHeader(404)
}

////////////////
//Internal API//
////////////////

func RestTransPkg(
    res http.ResponseWriter,
    req *http.Request,
    ram martini.Params,
    env *APIEnv,
) {
    pkg := PackageCreate(ram["pkgid"],env)
    err := pkg.Get()
    if err == nil {
	res.Header().Set("X-Accel-Redirect","/internal" + pkg.Export())
	res.WriteHeader(307)
    } else if _,ok := err.(ErrPackageMiss); ok {
	PackageClean(pkg.pkgid,env)
	res.WriteHeader(404)
    } else {
	res.WriteHeader(404)
    }
}
<commit_msg>Fix transport redirect bug, use ServeFile directly<commit_after>package main

import (
    "os"
    "io"
    "net/http"
    "github.com/go-martini/martini"
    "github.com/martini-contrib/render"
//  "github.com/garyburd/redigo/redis"
)

////////////////
//External API//
////////////////

func RestAddPkg(ren render.Render,req *http.Request,env *APIEnv) {
    ret := map[string]interface{}{"error":""}
    defer ren.JSON(200,ret)

    pkglen := req.ContentLength
    if pkglen <= 0 {
	ret["error"] = "EINVAL"
	return
    }
    if pkglen > PKG_MAXSIZE {
	ret["error"] = "ELIMIT"
	return
    }

    pkgid := PackageGenID()
    pkgfpath := STORAGE_PATH + "/package/" + pkgid + ".tar.xz"
    pkgfile,err := os.Create(pkgfpath)
    if err != nil {
	ret["error"] = "EIO"
	return
    }

    retlen,err := io.CopyN(pkgfile,req.Body,pkglen)
    pkgfile.Close()
    if retlen != pkglen || err != nil {
	os.Remove(pkgfpath)
	ret["error"] = "EIO"
	return
    }

    pkg := PackageCreate(pkgid,env)
    if pkg.Import(pkgfpath) != nil {
	os.Remove(pkgfpath)
	ret["error"] = "EINVAL"
	return
    }

    ret["pkgid"] = pkg.pkgid
}
func RestGetPkg(
    res http.ResponseWriter,
    req *http.Request,
    ram martini.Params,
    env *APIEnv,
) {
    pkg := PackageCreate(ram["pkgid"],env)
    err := pkg.Get()
    if err == nil && pkg.apiid == env.apiid {
	res.Header().Set("X-Accel-Redirect","/internal" + pkg.Export())
	res.WriteHeader(307)
	return
    } else if _,ok := err.(ErrPackageMiss); ok {
	if pkg.Transport() == nil {
	    res.Header().Set("X-Accel-Redirect","/internal" + pkg.Export())
	    res.WriteHeader(307)
	    return
	}
    }

    res.WriteHeader(404)
}

////////////////
//Internal API//
////////////////

func RestTransPkg(
    res http.ResponseWriter,
    req *http.Request,
    ram martini.Params,
    env *APIEnv,
) {
    pkg := PackageCreate(ram["pkgid"],env)
    err := pkg.Get()
    if err == nil {
	http.ServeFile(res,req,STORAGE_PATH + pkg.Export())
    } else if _,ok := err.(ErrPackageMiss); ok {
	PackageClean(pkg.pkgid,env)
	res.WriteHeader(404)
    } else {
	res.WriteHeader(404)
    }
}
<|endoftext|>