<commit_before>// Copyright 2015 The Serulian Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package es5

// Note: nativenew is based on http://www.bennadel.com/blog/2291-invoking-a-native-javascript-constructor-using-call-or-apply.htm

// runtimeTemplate contains all the necessary code for wrapping generated modules into a complete Serulian
// runtime bundle.
const runtimeTemplate = `
window.Serulian = (function($global) {
  var $g = {};
  var $t = {
    'cast': function(value, type) {
      // TODO: implement cast checking.
      return value
    },

    'nativenew': function(type) {
      return function () {
        var newInstance = Object.create(type.prototype);
        newInstance = type.apply(newInstance, arguments) || newInstance;
        return newInstance;
      };
    },

    'property': function(isExtension, getter, opt_setter) {
      var count = isExtension ? 2 : 1;
      var f = function() {
        if (arguments.length == count) {
          return opt_setter.apply(this, arguments);
        } else {
          return getter.apply(this, arguments);
        }
      };

      f.$property = true;
      return f;
    },

    'dynamicaccess': function(obj, name) {
      if (obj == null || obj[name] == null) {
        return null;
      }

      var value = obj[name];
      if (typeof value == 'function' && value.$property) {
        return $promise.wrap(function() {
          return value.apply(obj, arguments);
        });
      }

      return value
    },

    'nullcompare': function(first, second) {
      return first == null ? second : first;
    },

  	'sm': function(caller) {
  		return {
        resources: {},
  			current: 0,
  			next: caller,

        pushr: function(value, name) {
          this.resources[name] = value;
        },

        popr: function(names) {
          var promises = [];

          for (var i = 0; i < arguments.length; ++i) {
            var name = arguments[i];
            if (this.resources[name]) {
              promises.push(this.resources[name].Release());
              delete this.resources[name];
            }
          }

          if (promises.length > 0) {
            return $promise.all(promises);
          } else {
            return $promise.resolve(null);
          }
        },

        popall: function() {
          for (var name in this.resources) {
            if (this.resources.hasOwnProperty(name)) {
              this.resources[name].Release();
            }
          }
        }
  		};
  	}
  };

  var $promise = {
  	'build': function(statemachine) {
  		return new Promise(function(resolve, reject) {
        statemachine.resolve = function(value) {
          statemachine.popall();
          statemachine.current = -1;
          resolve(value);
        };

        statemachine.reject = function(value) {
          statemachine.popall();
          statemachine.current = -1;
          reject(value);
        };

  			var continueFunc = function() {
  				if (statemachine.current < 0) {
  					return;
  				}

  				statemachine.next(continueFunc);
			  };

  			continueFunc();
        if (statemachine.current < 0) {
          statemachine.resolve(null);
        }
  		});
  	},

  	'all': function(promises) {
  		return Promise.all(promises);
  	},

  	'empty': function() {
  		return new Promise(function() {
  			resolve();
  		});
  	},

    'resolve': function(value) {
      return Promise.resolve(value);
    },

  	'wrap': function(func) {
  		return Promise.resolve(func());
  	},

    'translate': function(prom) {
       return {
          'then': function() {
             return prom.Then.apply(prom, arguments);
          },
          'catch': function() {
             return prom.Catch.apply(prom, arguments);
          }
       };
    }
  };

  var moduleInits = [];

  var $module = function(name, creator) {
  	var module = {};

    var parts = name.split('.');
    var current = $g;
    for (var i = 0; i < parts.length - 1; ++i) {
      if (!current[parts[i]]) {
        current[parts[i]] = {};
      }
      current = current[parts[i]]
    }

    current[parts[parts.length - 1]] = module;

  	module.$init = function(cpromise) {
  	  moduleInits.push(cpromise);
  	};

    module.$newtypebuilder = function(kind) {
      return function(name, hasGenerics, creator) {
        if (hasGenerics) {
          module[name] = function(genericargs) {
            var tpe = function() {};
            creator.apply(tpe, arguments);
            return tpe;
          };
        } else {
          var tpe = function() {};
          creator.call(tpe);
          module[name] = tpe;
        }
      };
    };

  	module.$class = module.$newtypebuilder('class');
  	module.$interface = module.$newtypebuilder('interface');

    module.$type = function(name, creator) {
      var cls = function() {};
      creator.call(cls);
      module[name] = cls;
    };

  	creator.call(module)
  };

  {{ range $idx, $kv := .Iter }}
  	{{ $kv.Value }}
  {{ end }}

  return $promise.all(moduleInits).then(function() {
  	return $g;
  });
})(window)
`
<commit_msg>Give generated types a nice name<commit_after>// Copyright 2015 The Serulian Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package es5

// Note: nativenew is based on http://www.bennadel.com/blog/2291-invoking-a-native-javascript-constructor-using-call-or-apply.htm

// runtimeTemplate contains all the necessary code for wrapping generated modules into a complete Serulian
// runtime bundle.
const runtimeTemplate = `
window.Serulian = (function($global) {
  var $g = {};
  var $t = {
    'cast': function(value, type) {
      // TODO: implement cast checking.
      return value
    },

    'nativenew': function(type) {
      return function () {
        var newInstance = Object.create(type.prototype);
        newInstance = type.apply(newInstance, arguments) || newInstance;
        return newInstance;
      };
    },

    'property': function(isExtension, getter, opt_setter) {
      var count = isExtension ? 2 : 1;
      var f = function() {
        if (arguments.length == count) {
          return opt_setter.apply(this, arguments);
        } else {
          return getter.apply(this, arguments);
        }
      };

      f.$property = true;
      return f;
    },

    'dynamicaccess': function(obj, name) {
      if (obj == null || obj[name] == null) {
        return null;
      }

      var value = obj[name];
      if (typeof value == 'function' && value.$property) {
        return $promise.wrap(function() {
          return value.apply(obj, arguments);
        });
      }

      return value
    },

    'nullcompare': function(first, second) {
      return first == null ? second : first;
    },

  	'sm': function(caller) {
  		return {
        resources: {},
  			current: 0,
  			next: caller,

        pushr: function(value, name) {
          this.resources[name] = value;
        },

        popr: function(names) {
          var promises = [];

          for (var i = 0; i < arguments.length; ++i) {
            var name = arguments[i];
            if (this.resources[name]) {
              promises.push(this.resources[name].Release());
              delete this.resources[name];
            }
          }

          if (promises.length > 0) {
            return $promise.all(promises);
          } else {
            return $promise.resolve(null);
          }
        },

        popall: function() {
          for (var name in this.resources) {
            if (this.resources.hasOwnProperty(name)) {
              this.resources[name].Release();
            }
          }
        }
  		};
  	}
  };

  var $promise = {
  	'build': function(statemachine) {
  		return new Promise(function(resolve, reject) {
        statemachine.resolve = function(value) {
          statemachine.popall();
          statemachine.current = -1;
          resolve(value);
        };

        statemachine.reject = function(value) {
          statemachine.popall();
          statemachine.current = -1;
          reject(value);
        };

  			var continueFunc = function() {
  				if (statemachine.current < 0) {
  					return;
  				}

  				statemachine.next(continueFunc);
			  };

  			continueFunc();
        if (statemachine.current < 0) {
          statemachine.resolve(null);
        }
  		});
  	},

  	'all': function(promises) {
  		return Promise.all(promises);
  	},

  	'empty': function() {
  		return new Promise(function() {
  			resolve();
  		});
  	},

    'resolve': function(value) {
      return Promise.resolve(value);
    },

  	'wrap': function(func) {
  		return Promise.resolve(func());
  	},

    'translate': function(prom) {
       return {
          'then': function() {
             return prom.Then.apply(prom, arguments);
          },
          'catch': function() {
             return prom.Catch.apply(prom, arguments);
          }
       };
    }
  };

  var moduleInits = [];

  var $module = function(name, creator) {
  	var module = {};

    var parts = name.split('.');
    var current = $g;
    for (var i = 0; i < parts.length - 1; ++i) {
      if (!current[parts[i]]) {
        current[parts[i]] = {};
      }
      current = current[parts[i]]
    }

    current[parts[parts.length - 1]] = module;

  	module.$init = function(cpromise) {
  	  moduleInits.push(cpromise);
  	};

    module.$newtypebuilder = function(kind) {
      return function(name, hasGenerics, creator) {
        if (hasGenerics) {
          module[name] = function(genericargs) {
            var tpe = function() {};
            creator.apply(tpe, arguments);
            return tpe;
          };
        } else {
          var tpe = new Function("return function " + name + "() {};")();
          creator.call(tpe);
          module[name] = tpe;
        }
      };
    };

  	module.$class = module.$newtypebuilder('class');
  	module.$interface = module.$newtypebuilder('interface');

    module.$type = function(name, creator) {
      var cls = function() {};
      creator.call(cls);
      module[name] = cls;
    };

  	creator.call(module)
  };

  {{ range $idx, $kv := .Iter }}
  	{{ $kv.Value }}
  {{ end }}

  return $promise.all(moduleInits).then(function() {
  	return $g;
  });
})(window)
`
<|endoftext|>


<commit_before>package kasper

import (
	"log"

	"github.com/Shopify/sarama"
)

type partitionProcessor struct {
	topicProcessor                  *TopicProcessor
	coordinator                     Coordinator
	consumer                        sarama.Consumer
	partitionConsumers              []sarama.PartitionConsumer
	offsetManagers                  map[string]sarama.PartitionOffsetManager
	messageProcessor                MessageProcessor
	inputTopics                     []string
	partition                       int
	inFlightMessageGroups           map[string][]*inFlightMessageGroup
	messageProcessorRequestedCommit bool
}

func (pp *partitionProcessor) consumerMessageChannels() []<-chan *sarama.ConsumerMessage {
	chans := make([]<-chan *sarama.ConsumerMessage, len(pp.partitionConsumers))
	for i, consumer := range pp.partitionConsumers {
		chans[i] = consumer.Messages()
	}
	return chans
}

func newPartitionProcessor(tp *TopicProcessor, mp MessageProcessor, partition int) *partitionProcessor {
	consumer, err := sarama.NewConsumerFromClient(tp.client)
	if err != nil {
		log.Fatal(err)
	}
	partitionConsumers := make([]sarama.PartitionConsumer, len(tp.inputTopics))
	partitionOffsetManagers := make(map[string]sarama.PartitionOffsetManager)
	for i, topic := range tp.inputTopics {
		pom, err := tp.offsetManager.ManagePartition(topic, int32(partition))
		if err != nil {
			log.Fatal(err)
		}
		newestOffset, err := tp.client.GetOffset(topic, int32(partition), sarama.OffsetNewest)
		if err != nil {
			log.Fatal(err)
		}
		nextOffset, _ := pom.NextOffset()
		if nextOffset > newestOffset {
			nextOffset = sarama.OffsetNewest
		}
		c, err := consumer.ConsumePartition(topic, int32(partition), nextOffset)
		if err != nil {
			log.Fatal(err)
		}
		partitionConsumers[i] = c
		partitionOffsetManagers[topic] = pom
	}
	pp := &partitionProcessor{
		tp,
		nil,
		consumer,
		partitionConsumers,
		partitionOffsetManagers,
		mp,
		tp.inputTopics,
		partition,
		make(map[string][]*inFlightMessageGroup),
		false,
	}
	pp.coordinator = &partitionProcessorCoordinator{pp}
	return pp
}

func (pp *partitionProcessor) process(consumerMessage *sarama.ConsumerMessage) ([]*sarama.ProducerMessage, bool) {
	topicSerde, ok := pp.topicProcessor.config.TopicSerdes[consumerMessage.Topic]
	if !ok {
		log.Fatalf("Could not find Serde for topic '%s'", consumerMessage.Topic)
	}
	incomingMessage := IncomingMessage{
		Topic:     consumerMessage.Topic,
		Partition: int(consumerMessage.Partition),
		Offset:    consumerMessage.Offset,
		Key:       topicSerde.KeySerde.Deserialize(consumerMessage.Key),
		Value:     topicSerde.ValueSerde.Deserialize(consumerMessage.Value),
		Timestamp: consumerMessage.Timestamp,
	}
	sender := newSender(pp, &incomingMessage)
	pp.messageProcessorRequestedCommit = false
	pp.messageProcessor.Process(incomingMessage, sender, pp.coordinator)
	inFlightMessageGroup := sender.createInFlightMessageGroup()
	pp.inFlightMessageGroups[consumerMessage.Topic] = append(
		pp.inFlightMessageGroups[consumerMessage.Topic],
		inFlightMessageGroup,
	)
	return sender.producerMessages, pp.messageProcessorRequestedCommit
}

func (pp *partitionProcessor) onProcessCompleted() {
	pp.pruneInFlightMessageGroups()
}

func (pp *partitionProcessor) pruneInFlightMessageGroups() {
	for _, topic := range pp.topicProcessor.inputTopics {
		pp.pruneInFlightMessageGroupsForTopic(topic)
	}
	pp.countInFlightMessages()
	pp.countMessagesBehindHighWaterMark()
}

func (pp *partitionProcessor) countInFlightMessages() {
	partition := string(pp.partition)
	for _, topic := range pp.topicProcessor.inputTopics {
		var count int
		for _, groups := range pp.inFlightMessageGroups[topic] {
			count += len(groups.inFlightMessages)
		}
		pp.topicProcessor.inFlightMessagesCount.Set(float64(count), topic, partition)
	}
}

func (pp *partitionProcessor) countMessagesBehindHighWaterMark() {
	partition := string(pp.partition)
	highWaterMarks := pp.consumer.HighWaterMarks()
	for _, topic := range pp.topicProcessor.inputTopics {
		offsetManager := pp.offsetManagers[topic]
		currentOffset, _ := offsetManager.NextOffset()
		highWaterMark := highWaterMarks[topic][int32(pp.partition)]
		pp.topicProcessor.messagesBehindHighWaterMark.Set(float64(highWaterMark - currentOffset), topic, partition)
	}
}

func (pp *partitionProcessor) pruneInFlightMessageGroupsForTopic(topic string) {
	for len(pp.inFlightMessageGroups[topic]) > 1 {
		headGroup := pp.inFlightMessageGroups[topic][0]
		nextGroup := pp.inFlightMessageGroups[topic][1]
		if !headGroup.allAcksAreTrue() || !nextGroup.allAcksAreTrue() {
			break
		}
		pp.inFlightMessageGroups[topic] = pp.inFlightMessageGroups[topic][1:]
	}
}

func (pp *partitionProcessor) isReadyForMessage(msg *sarama.ConsumerMessage) bool {
	maxGroups := pp.topicProcessor.config.Config.MaxInFlightMessageGroups
	return len(pp.inFlightMessageGroups[msg.Topic]) <= maxGroups
}

func (pp *partitionProcessor) onMarkOffsetsTick() {
	for _, topic := range pp.topicProcessor.inputTopics {
		pp.onMarkOffsetsTickForTopic(topic)
	}
}

func (pp *partitionProcessor) onMarkOffsetsTickForTopic(topic string) {
	var offset int64 = -1
	for len(pp.inFlightMessageGroups[topic]) > 0 {
		group := pp.inFlightMessageGroups[topic][0]
		if !group.allAcksAreTrue() {
			break
		}
		offset = group.incomingMessage.Offset
		pp.inFlightMessageGroups[topic] = pp.inFlightMessageGroups[topic][1:]
	}
	if offset != -1 {
		offsetManager := pp.offsetManagers[topic]
		offsetManager.MarkOffset(offset+1, "")
	}
}

func (pp *partitionProcessor) onProducerAck(sentMessage *sarama.ProducerMessage) {
	incomingMessage := sentMessage.Metadata.(*IncomingMessage)
	foundGroup := false
	for _, group := range pp.inFlightMessageGroups[incomingMessage.Topic] {
		if group.incomingMessage == incomingMessage {
			foundGroup = true
			foundMsg := false
			for _, inFlightMessage := range group.inFlightMessages {
				if inFlightMessage.msg == sentMessage {
					foundMsg = true
					inFlightMessage.ack = true
					break
				}
			}
			if !foundMsg {
				log.Fatal("Could not find producer message in inFlightMessageGroups")
			}
			break
		}
	}
	if !foundGroup {
		log.Fatal("Could not find group in inFlightMessageGroups")
	}
}

func (pp *partitionProcessor) onShutdown() {
	var err error
	for _, pom := range pp.offsetManagers {
		err = pom.Close()
		if err != nil {
			log.Printf("Cannot close offset manager: %s", err)
		}

	}
	for _, pc := range pp.partitionConsumers {
		err = pc.Close()
		if err != nil {
			log.Printf("Cannot close partition consumer: %s", err)
		}
	}
	err = pp.consumer.Close()
	if err != nil {
		log.Fatal(err)
	}
}

func (pp *partitionProcessor) isReadyToCommit() bool {
	pp.pruneInFlightMessageGroups()
	for _, topic := range pp.inputTopics {
		if len(pp.inFlightMessageGroups[topic]) == 0 {
			continue
		} else if len(pp.inFlightMessageGroups[topic]) > 1 {
			return false
		} else {
			group := pp.inFlightMessageGroups[topic][0]
			if !group.allAcksAreTrue() {
				return false
			}
		}
	}
	return true
}

func (pp *partitionProcessor) commit() {
	for _, topic := range pp.inputTopics {
		if len(pp.inFlightMessageGroups[topic]) == 0 {
			continue
		}
		group := pp.inFlightMessageGroups[topic][0]
		offset := group.incomingMessage.Offset
		offsetManager := pp.offsetManagers[topic]
		offsetManager.MarkOffset(offset+1, "")
	}
}
<commit_msg>Fix buggy casting<commit_after>package kasper

import (
	"log"

	"github.com/Shopify/sarama"
	"strconv"
)

type partitionProcessor struct {
	topicProcessor                  *TopicProcessor
	coordinator                     Coordinator
	consumer                        sarama.Consumer
	partitionConsumers              []sarama.PartitionConsumer
	offsetManagers                  map[string]sarama.PartitionOffsetManager
	messageProcessor                MessageProcessor
	inputTopics                     []string
	partition                       int
	inFlightMessageGroups           map[string][]*inFlightMessageGroup
	messageProcessorRequestedCommit bool
}

func (pp *partitionProcessor) consumerMessageChannels() []<-chan *sarama.ConsumerMessage {
	chans := make([]<-chan *sarama.ConsumerMessage, len(pp.partitionConsumers))
	for i, consumer := range pp.partitionConsumers {
		chans[i] = consumer.Messages()
	}
	return chans
}

func newPartitionProcessor(tp *TopicProcessor, mp MessageProcessor, partition int) *partitionProcessor {
	consumer, err := sarama.NewConsumerFromClient(tp.client)
	if err != nil {
		log.Fatal(err)
	}
	partitionConsumers := make([]sarama.PartitionConsumer, len(tp.inputTopics))
	partitionOffsetManagers := make(map[string]sarama.PartitionOffsetManager)
	for i, topic := range tp.inputTopics {
		pom, err := tp.offsetManager.ManagePartition(topic, int32(partition))
		if err != nil {
			log.Fatal(err)
		}
		newestOffset, err := tp.client.GetOffset(topic, int32(partition), sarama.OffsetNewest)
		if err != nil {
			log.Fatal(err)
		}
		nextOffset, _ := pom.NextOffset()
		if nextOffset > newestOffset {
			nextOffset = sarama.OffsetNewest
		}
		c, err := consumer.ConsumePartition(topic, int32(partition), nextOffset)
		if err != nil {
			log.Fatal(err)
		}
		partitionConsumers[i] = c
		partitionOffsetManagers[topic] = pom
	}
	pp := &partitionProcessor{
		tp,
		nil,
		consumer,
		partitionConsumers,
		partitionOffsetManagers,
		mp,
		tp.inputTopics,
		partition,
		make(map[string][]*inFlightMessageGroup),
		false,
	}
	pp.coordinator = &partitionProcessorCoordinator{pp}
	return pp
}

func (pp *partitionProcessor) process(consumerMessage *sarama.ConsumerMessage) ([]*sarama.ProducerMessage, bool) {
	topicSerde, ok := pp.topicProcessor.config.TopicSerdes[consumerMessage.Topic]
	if !ok {
		log.Fatalf("Could not find Serde for topic '%s'", consumerMessage.Topic)
	}
	incomingMessage := IncomingMessage{
		Topic:     consumerMessage.Topic,
		Partition: int(consumerMessage.Partition),
		Offset:    consumerMessage.Offset,
		Key:       topicSerde.KeySerde.Deserialize(consumerMessage.Key),
		Value:     topicSerde.ValueSerde.Deserialize(consumerMessage.Value),
		Timestamp: consumerMessage.Timestamp,
	}
	sender := newSender(pp, &incomingMessage)
	pp.messageProcessorRequestedCommit = false
	pp.messageProcessor.Process(incomingMessage, sender, pp.coordinator)
	inFlightMessageGroup := sender.createInFlightMessageGroup()
	pp.inFlightMessageGroups[consumerMessage.Topic] = append(
		pp.inFlightMessageGroups[consumerMessage.Topic],
		inFlightMessageGroup,
	)
	return sender.producerMessages, pp.messageProcessorRequestedCommit
}

func (pp *partitionProcessor) onProcessCompleted() {
	pp.pruneInFlightMessageGroups()
}

func (pp *partitionProcessor) pruneInFlightMessageGroups() {
	for _, topic := range pp.topicProcessor.inputTopics {
		pp.pruneInFlightMessageGroupsForTopic(topic)
	}
	pp.countInFlightMessages()
	pp.countMessagesBehindHighWaterMark()
}

func (pp *partitionProcessor) countInFlightMessages() {
	partition := strconv.Itoa(pp.partition)
	for _, topic := range pp.topicProcessor.inputTopics {
		var count int
		for _, groups := range pp.inFlightMessageGroups[topic] {
			count += len(groups.inFlightMessages)
		}
		pp.topicProcessor.inFlightMessagesCount.Set(float64(count), topic, partition)
	}
}

func (pp *partitionProcessor) countMessagesBehindHighWaterMark() {
	partition := strconv.Itoa(pp.partition)
	highWaterMarks := pp.consumer.HighWaterMarks()
	for _, topic := range pp.topicProcessor.inputTopics {
		offsetManager := pp.offsetManagers[topic]
		currentOffset, _ := offsetManager.NextOffset()
		highWaterMark := highWaterMarks[topic][int32(pp.partition)]
		pp.topicProcessor.messagesBehindHighWaterMark.Set(float64(highWaterMark - currentOffset), topic, partition)
	}
}

func (pp *partitionProcessor) pruneInFlightMessageGroupsForTopic(topic string) {
	for len(pp.inFlightMessageGroups[topic]) > 1 {
		headGroup := pp.inFlightMessageGroups[topic][0]
		nextGroup := pp.inFlightMessageGroups[topic][1]
		if !headGroup.allAcksAreTrue() || !nextGroup.allAcksAreTrue() {
			break
		}
		pp.inFlightMessageGroups[topic] = pp.inFlightMessageGroups[topic][1:]
	}
}

func (pp *partitionProcessor) isReadyForMessage(msg *sarama.ConsumerMessage) bool {
	maxGroups := pp.topicProcessor.config.Config.MaxInFlightMessageGroups
	return len(pp.inFlightMessageGroups[msg.Topic]) <= maxGroups
}

func (pp *partitionProcessor) onMarkOffsetsTick() {
	for _, topic := range pp.topicProcessor.inputTopics {
		pp.onMarkOffsetsTickForTopic(topic)
	}
}

func (pp *partitionProcessor) onMarkOffsetsTickForTopic(topic string) {
	var offset int64 = -1
	for len(pp.inFlightMessageGroups[topic]) > 0 {
		group := pp.inFlightMessageGroups[topic][0]
		if !group.allAcksAreTrue() {
			break
		}
		offset = group.incomingMessage.Offset
		pp.inFlightMessageGroups[topic] = pp.inFlightMessageGroups[topic][1:]
	}
	if offset != -1 {
		offsetManager := pp.offsetManagers[topic]
		offsetManager.MarkOffset(offset+1, "")
	}
}

func (pp *partitionProcessor) onProducerAck(sentMessage *sarama.ProducerMessage) {
	incomingMessage := sentMessage.Metadata.(*IncomingMessage)
	foundGroup := false
	for _, group := range pp.inFlightMessageGroups[incomingMessage.Topic] {
		if group.incomingMessage == incomingMessage {
			foundGroup = true
			foundMsg := false
			for _, inFlightMessage := range group.inFlightMessages {
				if inFlightMessage.msg == sentMessage {
					foundMsg = true
					inFlightMessage.ack = true
					break
				}
			}
			if !foundMsg {
				log.Fatal("Could not find producer message in inFlightMessageGroups")
			}
			break
		}
	}
	if !foundGroup {
		log.Fatal("Could not find group in inFlightMessageGroups")
	}
}

func (pp *partitionProcessor) onShutdown() {
	var err error
	for _, pom := range pp.offsetManagers {
		err = pom.Close()
		if err != nil {
			log.Printf("Cannot close offset manager: %s", err)
		}

	}
	for _, pc := range pp.partitionConsumers {
		err = pc.Close()
		if err != nil {
			log.Printf("Cannot close partition consumer: %s", err)
		}
	}
	err = pp.consumer.Close()
	if err != nil {
		log.Fatal(err)
	}
}

func (pp *partitionProcessor) isReadyToCommit() bool {
	pp.pruneInFlightMessageGroups()
	for _, topic := range pp.inputTopics {
		if len(pp.inFlightMessageGroups[topic]) == 0 {
			continue
		} else if len(pp.inFlightMessageGroups[topic]) > 1 {
			return false
		} else {
			group := pp.inFlightMessageGroups[topic][0]
			if !group.allAcksAreTrue() {
				return false
			}
		}
	}
	return true
}

func (pp *partitionProcessor) commit() {
	for _, topic := range pp.inputTopics {
		if len(pp.inFlightMessageGroups[topic]) == 0 {
			continue
		}
		group := pp.inFlightMessageGroups[topic][0]
		offset := group.incomingMessage.Offset
		offsetManager := pp.offsetManagers[topic]
		offsetManager.MarkOffset(offset+1, "")
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 The golibpcap Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build linux,!safe,!appengine

package pkt

/*
#include "../pcap.h"
#include <net/ethernet.h>
#include <netinet/if_ether.h>
#include <netinet/in.h>

#include <netinet/ip.h>
#include <netinet/tcp.h>
*/
import "C"
import (
	"reflect"
	"time"
	"unsafe"
)

// The Packet struct is a wrapper for the pcap_pkthdr struct in <pcap.h>.
type Packet struct {
	Time    time.Time      // time stamp from the nic
	Caplen  uint32         // length of portion present
	Len     uint32         // length this packet (off wire)
	Headers []Hdr          // Go wrappers for C pkt headers
	buf     unsafe.Pointer // packet data (*C.u_char)
}

// NewPacket returns a parsed and decoded Packet.
// pkthdr_ptr should be a *C.struct_pcap_pkthdr
// buf_ptr should be a *C.u_char
func NewPacket(pkthdr_ptr unsafe.Pointer, buf_ptr unsafe.Pointer) *Packet {
	pkthdr := *(*C.struct_pcap_pkthdr)(pkthdr_ptr)

	p := &Packet{
		Time:    time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000),
		Caplen:  uint32(pkthdr.caplen),
		Len:     uint32(pkthdr.len),
		Headers: make([]Hdr, 3),
		buf:     buf_ptr,
	}
	p.decode()
	return p
}

// Decode decodes the headers of a Packet.
func (p *Packet) decode() {
	ethHdr, buf := NewEthHdr(p.buf)
	p.Headers[LinkLayer] = ethHdr

	switch ethHdr.EtherType {
	case C.ETHERTYPE_IP, 0:
		p.Headers[NetworkLayer], buf = NewIpHdr(buf)
	case C.ETHERTYPE_IPV6:
		p.Headers[NetworkLayer], buf = NewIp6Hdr(buf)
	case C.ETHERTYPE_ARP:
		//TODO(gavaletz) ARP
		return
	default:
		return
	}

	switch p.Headers[NetworkLayer].(InetProtoHdr).Proto() {
	case C.IPPROTO_TCP:
		p.Headers[TransportLayer], _ = NewTcpHdr(buf)
	case C.IPPROTO_UDP:
		p.Headers[TransportLayer], _ = NewUdpHdr(buf)
		return
	case C.IPPROTO_ICMP:
		//TODO(gavaletz) ICMP
		return
	default:
		return
	}
}

type TcpPacket struct {
	DstAddr   uint32
	SrcAddr   uint32
	AckSeq    uint32
	Seq       uint32
	Source    uint16
	Dest      uint16
	Flags     uint16
	Payload   []byte
	Timestamp time.Time
	IsRequest bool
}

func (this *TcpPacket) Save() {
	var dcopy = make([]byte, len(this.Payload))
	copy(dcopy, this.Payload)
	this.Payload = dcopy
}

func NewPacket2(pkthdr_ptr unsafe.Pointer, buf_ptr unsafe.Pointer) *TcpPacket {
	pkthdr := *(*C.struct_pcap_pkthdr)(pkthdr_ptr)

	// unwrap ethernet packet
	var ethhdr = (*C.struct_ether_header)(buf_ptr)
	var ethtype = uint16(ethhdr.ether_type)
	// we are assuming little endian arch
	ethtype = (ethtype>>8 | ethtype&uint16(0x00ff)<<8)

	if ethtype == 0 {
		// The "cooked" headers have an extra two bytes.
		buf_ptr = unsafe.Pointer(uintptr(buf_ptr) + uintptr(C.ETHER_HDR_LEN) + uintptr(2))
	} else {
		buf_ptr = unsafe.Pointer(uintptr(buf_ptr) + uintptr(C.ETHER_HDR_LEN))
	}

	if ethtype != C.ETHERTYPE_IP && ethtype != 0 {
		return nil
	}

	// unwrap ip packet
	var iphdr = (*C.struct_iphdr)(buf_ptr)
	var iphdrlen = *(*byte)(buf_ptr) & 0x0F

	var paylen = uint16(iphdr.tot_len)
	// we are assuming little endian arch
	paylen = (paylen>>8 | paylen&uint16(0x00ff)<<8) - uint16(iphdrlen*4)

	buf_ptr = unsafe.Pointer(uintptr(buf_ptr) + uintptr(iphdrlen*4))

	if uint8(iphdr.protocol) != C.IPPROTO_TCP {
		return nil
	}

	// unwrap tcp packet
	var tcphdr = (*C.struct_tcphdr)(buf_ptr)
	var dataoffset = *(*byte)(unsafe.Pointer(uintptr(buf_ptr) + uintptr(12))) >> 4

	packet := &TcpPacket{
		DstAddr:   uint32(iphdr.daddr),
		SrcAddr:   uint32(iphdr.saddr),
		AckSeq:    uint32(tcphdr.ack_seq),
		Seq:       uint32(tcphdr.seq),
		Source:    uint16(tcphdr.source),
		Dest:      uint16(tcphdr.dest),
		Flags:     *(*uint16)(unsafe.Pointer(uintptr(buf_ptr) + uintptr(12))),
		Timestamp: time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000),
		IsRequest: false,
	}

	sh := (*reflect.SliceHeader)((unsafe.Pointer(&packet.Payload)))
	sh.Cap = int(paylen - uint16(dataoffset*4))
	sh.Len = sh.Cap
	sh.Data = uintptr(unsafe.Pointer(uintptr(buf_ptr) + uintptr(dataoffset*4)))

	// Network to hosts. Right now we are assuming little endian cpu.
	packet.Flags = (packet.Flags>>8 | packet.Flags&uint16(0x00ff)<<8) & uint16(0x01FF)
	packet.Dest = packet.Dest>>8 | packet.Dest&uint16(0x00ff)<<8
	packet.Source = packet.Source>>8 | packet.Source&uint16(0x00ff)<<8
	packet.Seq = packet.Seq>>24 | packet.Seq&uint32(0x00ff0000)>>8 | packet.Seq&uint32(0x0000ff00)<<8 | packet.Seq<<24
	packet.AckSeq = packet.AckSeq>>24 | packet.AckSeq&uint32(0x00ff0000)>>8 | packet.AckSeq&uint32(0x0000ff00)<<8 | packet.AckSeq<<24
	return packet
}
<commit_msg>Only create new packets if the last one was saved<commit_after>// Copyright 2013 The golibpcap Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build linux,!safe,!appengine

package pkt

/*
#include "../pcap.h"
#include <net/ethernet.h>
#include <netinet/if_ether.h>
#include <netinet/in.h>

#include <netinet/ip.h>
#include <netinet/tcp.h>
*/
import "C"
import (
	"reflect"
	"time"
	"unsafe"
)

// The Packet struct is a wrapper for the pcap_pkthdr struct in <pcap.h>.
type Packet struct {
	Time    time.Time      // time stamp from the nic
	Caplen  uint32         // length of portion present
	Len     uint32         // length this packet (off wire)
	Headers []Hdr          // Go wrappers for C pkt headers
	buf     unsafe.Pointer // packet data (*C.u_char)
}

// NewPacket returns a parsed and decoded Packet.
// pkthdr_ptr should be a *C.struct_pcap_pkthdr
// buf_ptr should be a *C.u_char
func NewPacket(pkthdr_ptr unsafe.Pointer, buf_ptr unsafe.Pointer) *Packet {
	pkthdr := *(*C.struct_pcap_pkthdr)(pkthdr_ptr)

	p := &Packet{
		Time:    time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000),
		Caplen:  uint32(pkthdr.caplen),
		Len:     uint32(pkthdr.len),
		Headers: make([]Hdr, 3),
		buf:     buf_ptr,
	}
	p.decode()
	return p
}

// Decode decodes the headers of a Packet.
func (p *Packet) decode() {
	ethHdr, buf := NewEthHdr(p.buf)
	p.Headers[LinkLayer] = ethHdr

	switch ethHdr.EtherType {
	case C.ETHERTYPE_IP, 0:
		p.Headers[NetworkLayer], buf = NewIpHdr(buf)
	case C.ETHERTYPE_IPV6:
		p.Headers[NetworkLayer], buf = NewIp6Hdr(buf)
	case C.ETHERTYPE_ARP:
		//TODO(gavaletz) ARP
		return
	default:
		return
	}

	switch p.Headers[NetworkLayer].(InetProtoHdr).Proto() {
	case C.IPPROTO_TCP:
		p.Headers[TransportLayer], _ = NewTcpHdr(buf)
	case C.IPPROTO_UDP:
		p.Headers[TransportLayer], _ = NewUdpHdr(buf)
		return
	case C.IPPROTO_ICMP:
		//TODO(gavaletz) ICMP
		return
	default:
		return
	}
}

type TcpPacket struct {
	DstAddr   uint32
	SrcAddr   uint32
	AckSeq    uint32
	Seq       uint32
	Source    uint16
	Dest      uint16
	Flags     uint16
	Payload   []byte
	Timestamp time.Time
	IsRequest bool
	saved     bool
}

func (this *TcpPacket) Save() {
	var dcopy = make([]byte, len(this.Payload))
	copy(dcopy, this.Payload)
	this.Payload = dcopy
	this.saved = true
}

var packet *TcpPacket

func NewPacket2(pkthdr_ptr unsafe.Pointer, buf_ptr unsafe.Pointer) *TcpPacket {
	pkthdr := *(*C.struct_pcap_pkthdr)(pkthdr_ptr)

	// unwrap ethernet packet
	var ethhdr = (*C.struct_ether_header)(buf_ptr)
	var ethtype = uint16(ethhdr.ether_type)
	// we are assuming little endian arch
	ethtype = (ethtype>>8 | ethtype&uint16(0x00ff)<<8)

	if ethtype == 0 {
		// The "cooked" headers have an extra two bytes.
		buf_ptr = unsafe.Pointer(uintptr(buf_ptr) + uintptr(C.ETHER_HDR_LEN) + uintptr(2))
	} else {
		buf_ptr = unsafe.Pointer(uintptr(buf_ptr) + uintptr(C.ETHER_HDR_LEN))
	}

	if ethtype != C.ETHERTYPE_IP && ethtype != 0 {
		return nil
	}

	// unwrap ip packet
	var iphdr = (*C.struct_iphdr)(buf_ptr)
	var iphdrlen = *(*byte)(buf_ptr) & 0x0F

	var paylen = uint16(iphdr.tot_len)
	// we are assuming little endian arch
	paylen = (paylen>>8 | paylen&uint16(0x00ff)<<8) - uint16(iphdrlen*4)

	buf_ptr = unsafe.Pointer(uintptr(buf_ptr) + uintptr(iphdrlen*4))

	if uint8(iphdr.protocol) != C.IPPROTO_TCP {
		return nil
	}

	// unwrap tcp packet
	var tcphdr = (*C.struct_tcphdr)(buf_ptr)
	var dataoffset = *(*byte)(unsafe.Pointer(uintptr(buf_ptr) + uintptr(12))) >> 4

	if packet == nil || packet.saved {
		packet = &TcpPacket{}
	}

	packet.DstAddr = uint32(iphdr.daddr)
	packet.SrcAddr = uint32(iphdr.saddr)
	packet.AckSeq = uint32(tcphdr.ack_seq)
	packet.Seq = uint32(tcphdr.seq)
	packet.Source = uint16(tcphdr.source)
	packet.Dest = uint16(tcphdr.dest)
	packet.Flags = *(*uint16)(unsafe.Pointer(uintptr(buf_ptr) + uintptr(12)))
	packet.Timestamp = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000)
	packet.IsRequest = false

	sh := (*reflect.SliceHeader)((unsafe.Pointer(&packet.Payload)))
	sh.Cap = int(paylen - uint16(dataoffset*4))
	sh.Len = sh.Cap
	sh.Data = uintptr(unsafe.Pointer(uintptr(buf_ptr) + uintptr(dataoffset*4)))

	// Network to hosts. Right now we are assuming little endian cpu.
	packet.Flags = (packet.Flags>>8 | packet.Flags&uint16(0x00ff)<<8) & uint16(0x01FF)
	packet.Dest = packet.Dest>>8 | packet.Dest&uint16(0x00ff)<<8
	packet.Source = packet.Source>>8 | packet.Source&uint16(0x00ff)<<8
	packet.Seq = packet.Seq>>24 | packet.Seq&uint32(0x00ff0000)>>8 | packet.Seq&uint32(0x0000ff00)<<8 | packet.Seq<<24
	packet.AckSeq = packet.AckSeq>>24 | packet.AckSeq&uint32(0x00ff0000)>>8 | packet.AckSeq&uint32(0x0000ff00)<<8 | packet.AckSeq<<24
	return packet
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

   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 oci

import (
	specs "github.com/opencontainers/runtime-spec/specs-go"
)

func defaultMounts() []specs.Mount {
	return []specs.Mount{
		{
			Destination: "/dev",
			Type:        "devfs",
			Source:      "devfs",
			Options:     []string{},
		},
		{
			Destination: "/dev/fd",
			Type:        "fdescfs",
			Source:      "fdescfs",
			Options:     []string{},
		},
	}
}
<commit_msg>Add ruleset=4 option<commit_after>/*
   Copyright The containerd Authors.

   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 oci

import (
	specs "github.com/opencontainers/runtime-spec/specs-go"
)

func defaultMounts() []specs.Mount {
	return []specs.Mount{
		{
			Destination: "/dev",
			Type:        "devfs",
			Source:      "devfs",
			Options:     []string{"ruleset=4"},
		},
		{
			Destination: "/dev/fd",
			Type:        "fdescfs",
			Source:      "fdescfs",
			Options:     []string{},
		},
	}
}
<|endoftext|>


<commit_before>// vault is an implementation of a source using Hashicorp Vault to store the data
package vault

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"

	"github.com/hashicorp/vault/api"
	"github.com/pkg/errors"

	"github.com/MiLk/nsscache-go/cache"
)

type VaultSource struct {
	client *api.Client
	prefix string
}

type Option func(*VaultSource)

func Client(c *api.Client) Option {
	return func(s *VaultSource) { s.client = c }
}

func Prefix(p string) Option {
	return func(s *VaultSource) { s.prefix = p }
}

func NewSource(opts ...Option) (*VaultSource, error) {
	s := VaultSource{
		prefix: "nsscache",
	}

	for _, opt := range opts {
		opt(&s)
	}

	if s.client == nil {
		cl, err := api.NewClient(nil)
		if err != nil {
			return nil, err
		}
		s.client = cl
	}

	return &s, nil
}

func (s *VaultSource) Client() *api.Client {
	return s.client
}

func (s *VaultSource) list(name string, c *cache.Cache, createEntry func() cache.Entry) error {
	prefix := fmt.Sprintf("secret/%s/%s", s.prefix, name)
	sec, err := s.client.Logical().List(prefix)
	if err != nil {
		return errors.Wrap(err, "list from vault")
	}

	// No secret at that path
	if sec == nil {
		return nil
	}

	keys := sec.Data["keys"].([]interface{})
	for _, k := range keys {
		sec, err := s.client.Logical().Read(fmt.Sprintf("%s/%s", prefix, k))
		if err != nil {
			return errors.Wrap(err, "read from vault")
		}
		value := sec.Data["value"].(string)
		b := bytes.NewBufferString(value)
		b64 := base64.NewDecoder(base64.StdEncoding, b)
		e := createEntry()
		err = json.NewDecoder(b64).Decode(e)
		if err != nil {
			return errors.Wrap(err, "json decoding")
		}
		c.Add(e)
	}
	return nil
}

func (s *VaultSource) FillPasswdCache(c *cache.Cache) error {
	return s.list("passwd", c, func() cache.Entry {
		return &cache.PasswdEntry{}
	})
}

func (s *VaultSource) FillShadowCache(c *cache.Cache) error {
	return s.list("shadow", c, func() cache.Entry {
		return &cache.ShadowEntry{}
	})
}

func (s *VaultSource) FillGroupCache(c *cache.Cache) error {
	return s.list("group", c, func() cache.Entry {
		return &cache.GroupEntry{}
	})
}
<commit_msg>Fix vault calls for kv v2<commit_after>// vault is an implementation of a source using Hashicorp Vault to store the data
package vault

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"

	"github.com/hashicorp/vault/api"
	"github.com/pkg/errors"

	"github.com/MiLk/nsscache-go/cache"
)

type VaultSource struct {
	client    *api.Client
	prefix    string
	mountPath string
}

type Option func(*VaultSource)

func Client(c *api.Client) Option {
	return func(s *VaultSource) { s.client = c }
}

func Prefix(p string) Option {
	return func(s *VaultSource) { s.prefix = p }
}

func MountPath(m string) Option {
	return func(s *VaultSource) { s.mountPath = m }
}

func NewSource(opts ...Option) (*VaultSource, error) {
	s := VaultSource{
		prefix:    "nsscache",
		mountPath: "secret",
	}

	for _, opt := range opts {
		opt(&s)
	}

	if s.client == nil {
		cl, err := api.NewClient(nil)
		if err != nil {
			return nil, err
		}
		s.client = cl
	}

	return &s, nil
}

func (s *VaultSource) Client() *api.Client {
	return s.client
}

func (s *VaultSource) list(name string, c *cache.Cache, createEntry func() cache.Entry) error {
	prefix := fmt.Sprintf("%s/%s", s.prefix, name)
	sec, err := s.client.Logical().List(fmt.Sprintf("%s/metadata/%s", s.mountPath, prefix))
	if err != nil {
		return errors.Wrap(err, "list from vault")
	}

	// No secret at that path
	if sec == nil {
		return nil
	}

	keys := sec.Data["keys"].([]interface{})
	for _, k := range keys {
		sec, err := s.client.Logical().Read(fmt.Sprintf("%s/data/%s/%s", s.mountPath, prefix, k))
		if err != nil {
			return errors.Wrap(err, "read from vault")
		}
		value := sec.Data["data"].(map[string]interface{})["value"].(string)
		b := bytes.NewBufferString(value)
		b64 := base64.NewDecoder(base64.StdEncoding, b)
		e := createEntry()
		err = json.NewDecoder(b64).Decode(e)
		if err != nil {
			return errors.Wrap(err, "json decoding")
		}
		c.Add(e)
	}
	return nil
}

func (s *VaultSource) FillPasswdCache(c *cache.Cache) error {
	return s.list("passwd", c, func() cache.Entry {
		return &cache.PasswdEntry{}
	})
}

func (s *VaultSource) FillShadowCache(c *cache.Cache) error {
	return s.list("shadow", c, func() cache.Entry {
		return &cache.ShadowEntry{}
	})
}

func (s *VaultSource) FillGroupCache(c *cache.Cache) error {
	return s.list("group", c, func() cache.Entry {
		return &cache.GroupEntry{}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/paulvollmer/commenttags"
	"io/ioutil"
	"os"
)

const Version = "0.1.0"

func main() {
	flagVersion := flag.Bool("v", false, "print out the `version`")
	flagFormat := flag.String("f", "pretty", "set the output `format`. (pretty, json or json-pretty)")
	flagWrite := flag.String("w", "", "`write` to file")
	flagMaxFilesize := flag.Int64("m", 5000000, "the `maximum filesize` to process")
	flag.Parse()
	if *flagVersion {
		fmt.Println(Version)
		return
	}

	if len(os.Args) > 1 {
		// TODO: check if file or directory
		sourcePath := os.Args[len(os.Args)-1]
		fileResult, err := commenttags.ProcessFile(sourcePath)
		if err != nil {
			// fmt.Println("File Processing failed!", err)
			if err.Error() == "read "+sourcePath+": is a directory" {
				// fmt.Println("Try to read as Directory...")
				dirResult, errDir := commenttags.ProcessDirectory(sourcePath, *flagMaxFilesize)
				if errDir != nil {
					fmt.Println("Read Directory failed!", err)
					return
				}
				// dirResult PrettyPrint()
				// fmt.Println(dirResult)
				for _, v := range dirResult.Files {
					v.PrettyPrint()
				}
			}
			return
		}

		// format the result result
		formatted := ""
		switch *flagFormat {
		case "pretty":
			formatted = fileResult.Pretty()
			break
		case "json":
			out, _ := json.Marshal(fileResult)
			formatted = string(out)
			break
		case "json-pretty":
			out, _ := json.MarshalIndent(fileResult, "", "  ")
			formatted = string(out)
			break
		default:
			fmt.Println("Format not supported")
			break
		}

		// write or print out the result
		if *flagWrite != "" {
			ioutil.WriteFile(*flagWrite, []byte(formatted), 0777)
		} else {
			fmt.Println(formatted)
		}

	} else {
		fmt.Println("Missing Filepath, See the -h help out...")
	}
}
<commit_msg>Cli fixed format json-pretty issue<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/paulvollmer/commenttags"
	"io/ioutil"
	"os"
)

const (
	Version          = "0.1.0"
	SourceTypeFile   = 0
	SourceTypeFolder = 1
)

var (
	formatted  string
	sourceType int
)

func main() {
	flagVersion := flag.Bool("v", false, "print out the `version`")
	flagFormat := flag.String("f", "pretty", "set the output `format`. (pretty, json or json-pretty)")
	flagWrite := flag.String("w", "", "`write` to file")
	flagMaxFilesize := flag.Int64("m", 5000000, "the `maximum filesize` to process")
	flag.Parse()
	if *flagVersion {
		fmt.Println(Version)
		return
	}

	if len(os.Args) > 1 {
		//
		// process data
		//
		sourcePath := os.Args[len(os.Args)-1]
		folderResult := &commenttags.DirectoryData{}
		fileResult, err := commenttags.ProcessFile(sourcePath)
		if err != nil {
			// check if directory
			if err.Error() == "read "+sourcePath+": is a directory" {
				sourceType = SourceTypeFolder
				// Try to read as directory...
				folderResult, err = commenttags.ProcessDirectory(sourcePath, *flagMaxFilesize)
				if err != nil {
					fmt.Println("read directory failed!", err)
					os.Exit(1)
				}
			}
		}

		//
		// format the result result
		//
		switch *flagFormat {
		case "pretty":
			println("ok, pretty")
			if sourceType == SourceTypeFile {
				formatted = fileResult.Pretty()
			} else if sourceType == SourceTypeFolder {
				for _, v := range folderResult.Files {
					formatted += v.Pretty()
				}
			}
			break
		case "json":
			if sourceType == SourceTypeFile {
				out, _ := fileResult.JSON()
				formatted = string(out)
			} else if sourceType == SourceTypeFolder {
				tmp, _ := folderResult.JSON()
				formatted = string(tmp)
			}
			break
		case "json-pretty":
			if sourceType == SourceTypeFile {
				out, _ := json.MarshalIndent(fileResult, "", "  ")
				formatted = string(out)
			} else if sourceType == SourceTypeFolder {
				out, _ := json.MarshalIndent(folderResult, "", "  ")
				formatted = string(out)
			}
			break
		default:
			fmt.Printf("Format '%s' not supported! Choose between the following formats:\n", *flagFormat)
			fmt.Println("- pretty (default)")
			fmt.Println("- json")
			fmt.Println("- json-pretty")
			break
		}

		//
		// write or print out the result
		//
		if *flagWrite != "" {
			ioutil.WriteFile(*flagWrite, []byte(formatted), 0777)
		} else {
			fmt.Println(formatted)
		}

	} else {
		fmt.Println("Missing Filepath, See the -h help out...")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/guardian/rundmc/dadoo"

	"github.com/eapache/go-resiliency/retrier"
	"github.com/opencontainers/runc/libcontainer/system"

	cmsg "github.com/opencontainers/runc/libcontainer/utils"
)

const MaxSocketDirPathLength = 80

func main() {
	os.Exit(run())
}

func run() int {
	tty := flag.Bool("tty", false, "tty requested")
	socketDirPath := flag.String("socket-dir-path", "", "path to a dir in which to store console sockets")
	runcRoot := flag.String("runc-root", "", "root directory for storage of container state (this should be located in tmpfs)")
	flag.Parse()

	runtime := flag.Args()[1] // e.g. runc
	processStateDir := flag.Args()[2]
	containerId := flag.Args()[3]

	signals := make(chan os.Signal, 100)
	signal.Notify(signals, syscall.SIGCHLD)

	runcExitCodePipe := os.NewFile(3, "/proc/self/fd/3")
	logFile := fmt.Sprintf("/proc/%d/fd/4", os.Getpid())
	logFD := os.NewFile(4, "/proc/self/fd/4")
	syncPipe := os.NewFile(5, "/proc/self/fd/5")
	pidFilePath := filepath.Join(processStateDir, "pidfile")

	stdinR, stdoutW, stderrW, winsz := openPipes(processStateDir)
	defer func() {
		tryClose(stdinR, stdoutW, stderrW, winsz)
	}()

	syncPipe.Write([]byte{0})

	stdoutR, stderrR := openStdioKeepAlivePipes(processStateDir)
	defer func() {
		tryClose(stdoutR, stderrR)
	}()

	ioWg := &sync.WaitGroup{}
	var runcExecCmd *exec.Cmd
	runtimeArgs := []string{"-debug", "-log", logFile}
	if *runcRoot != "" {
		runtimeArgs = append(runtimeArgs, "-root", *runcRoot)
	}
	runtimeArgs = append(runtimeArgs, "exec", "-d", "-p", fmt.Sprintf("/proc/%d/fd/0", os.Getpid()), "-pid-file", pidFilePath)
	if *tty {
		if len(*socketDirPath) > MaxSocketDirPathLength {
			logAndExit(fmt.Sprintf("value for --socket-dir-path cannot exceed %d characters in length", MaxSocketDirPathLength))
		}
		ttySocketPath := setupTTYSocket(stdinR, stdoutW, winsz, pidFilePath, *socketDirPath, ioWg)
		runtimeArgs = append(runtimeArgs, "-tty", "-console-socket", ttySocketPath, containerId)
		runcExecCmd = exec.Command(runtime, runtimeArgs...)
	} else {
		runtimeArgs = append(runtimeArgs, containerId)
		runcExecCmd = exec.Command(runtime, runtimeArgs...)
		runcExecCmd.Stdin = stdinR
		runcExecCmd.Stdout = stdoutW
		runcExecCmd.Stderr = stderrW
	}

	// we need to be the subreaper so we can wait on the detached container process
	system.SetSubreaper(os.Getpid())

	if err := runcExecCmd.Start(); err != nil {
		runcExitCodePipe.Write([]byte{2})
		return 2
	}

	var status syscall.WaitStatus
	var rusage syscall.Rusage
	_, err := syscall.Wait4(runcExecCmd.Process.Pid, &status, 0, &rusage)
	check(err)    // Start succeeded but Wait4 failed, this can only be a programmer error
	logFD.Close() // No more logs from runc so close fd

	// also check that masterFD is received and streaming or whatevs
	runcExitCodePipe.Write([]byte{byte(status.ExitStatus())})
	if status.ExitStatus() != 0 {
		return 3 // nothing to wait for, container didn't launch
	}

	containerPid, err := parsePid(pidFilePath)
	check(err)

	return waitForContainerToExit(processStateDir, containerPid, signals, ioWg)
}

// If gdn server process dies, we need dadoo to keep stdout/err reader
// FDs so that Linux does not SIGPIPE the user process if it tries to use its end of
// these pipes.
func openStdioKeepAlivePipes(processStateDir string) (io.ReadCloser, io.ReadCloser) {
	keepStdoutAlive := openFifo(filepath.Join(processStateDir, "stdout"), os.O_RDONLY)
	keepStderrAlive := openFifo(filepath.Join(processStateDir, "stderr"), os.O_RDONLY)
	return keepStdoutAlive, keepStderrAlive
}

func waitForContainerToExit(processStateDir string, containerPid int, signals chan os.Signal, ioWg *sync.WaitGroup) (exitCode int) {
	for range signals {
		for {
			var status syscall.WaitStatus
			var rusage syscall.Rusage
			wpid, err := syscall.Wait4(-1, &status, syscall.WNOHANG, &rusage)
			if err != nil || wpid <= 0 {
				break // wait for next SIGCHLD
			}

			if wpid == containerPid {
				exitCode = status.ExitStatus()
				if status.Signaled() {
					exitCode = 128 + int(status.Signal())
				}

				ioWg.Wait() // wait for full output to be collected

				check(ioutil.WriteFile(filepath.Join(processStateDir, "exitcode"), []byte(strconv.Itoa(exitCode)), 0600))
				return exitCode
			}
		}
	}

	logAndExit("ran out of signals") // cant happen
	return 0                         // unreachable
}

func openPipes(processStateDir string) (io.ReadCloser, io.WriteCloser, io.WriteCloser, io.ReadWriteCloser) {
	stdin := openFifo(filepath.Join(processStateDir, "stdin"), os.O_RDONLY)
	stdout := openFifo(filepath.Join(processStateDir, "stdout"), os.O_WRONLY|os.O_APPEND)
	stderr := openFifo(filepath.Join(processStateDir, "stderr"), os.O_WRONLY|os.O_APPEND)
	winsz := openFifo(filepath.Join(processStateDir, "winsz"), os.O_RDWR)
	openFifo(filepath.Join(processStateDir, "exit"), os.O_RDWR) // open just so guardian can detect it being closed when we exit

	return stdin, stdout, stderr, winsz
}

func openFifo(path string, flags int) io.ReadWriteCloser {
	r, err := os.OpenFile(path, flags, 0600)
	if os.IsNotExist(err) {
		return nil
	}

	check(err)
	return r
}

func setupTTYSocket(stdin io.Reader, stdout io.Writer, winszFifo io.Reader, pidFilePath, sockDirBase string, ioWg *sync.WaitGroup) string {
	sockDir, err := ioutil.TempDir(sockDirBase, "")
	check(err)

	ttySockPath := filepath.Join(sockDir, "tty.sock")
	l, err := net.Listen("unix", ttySockPath)
	check(err)

	//go to the background and set master
	go func(ln net.Listener) (err error) {
		// if any of the following errors, it means runc has connected to the
		// socket, so it must've started, thus we might need to kill the process
		defer func() {
			if err != nil {
				killProcess(pidFilePath)
				check(err)
			}
		}()

		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer conn.Close()

		// Close ln, to allow for other instances to take over.
		ln.Close()

		// Get the fd of the connection.
		unixconn, ok := conn.(*net.UnixConn)
		if !ok {
			return
		}

		socket, err := unixconn.File()
		if err != nil {
			return
		}
		defer socket.Close()

		// Get the master file descriptor from runC.
		master, err := cmsg.RecvFd(socket)
		if err != nil {
			return
		}

		os.RemoveAll(sockDir)
		if err = setOnlcr(master); err != nil {
			return
		}
		streamProcess(master, stdin, stdout, winszFifo, ioWg)

		return
	}(l)

	return ttySockPath
}

func streamProcess(m *os.File, stdin io.Reader, stdout io.Writer, winszFifo io.Reader, ioWg *sync.WaitGroup) {
	ioWg.Add(1)
	go func() {
		defer ioWg.Done()
		io.Copy(stdout, m)
	}()

	go io.Copy(m, stdin)

	go func() {
		for {
			var winSize garden.WindowSize
			if err := json.NewDecoder(winszFifo).Decode(&winSize); err != nil {
				fmt.Printf("invalid winsz event: %s\n", err)
				continue // not much we can do here..
			}
			dadoo.SetWinSize(m, winSize)
		}
	}()
}

func killProcess(pidFilePath string) {
	pid, err := readPid(pidFilePath)
	if err == nil {
		syscall.Kill(pid, syscall.SIGKILL)
	}
}

func readPid(pidFilePath string) (int, error) {
	retrier := retrier.New(retrier.ConstantBackoff(20, 500*time.Millisecond), nil)
	var (
		pid int = -1
		err error
	)
	retrier.Run(func() error {
		pid, err = parsePid(pidFilePath)
		return err
	})

	return pid, err
}

func parsePid(pidFile string) (int, error) {
	b, err := ioutil.ReadFile(pidFile)
	if err != nil {
		return -1, err
	}

	var pid int
	if _, err := fmt.Sscanf(string(b), "%d", &pid); err != nil {
		return -1, err
	}

	return pid, nil
}

func logAndExit(msg string) {
	fmt.Println(msg)
	os.Exit(2)
}

func check(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(2)
	}
}

func tryClose(closers ...io.Closer) {
	for _, closer := range closers {
		if closer != nil {
			closer.Close()
		}
	}
}

// setOnlcr copied from runc
// https://github.com/cloudfoundry-incubator/runc/blob/02ec89829b24dfce45bb207d2344e0e6d078a93c/libcontainer/console_linux.go#L144-L160
func setOnlcr(terminal *os.File) error {
	var termios syscall.Termios

	if err := ioctl(terminal.Fd(), syscall.TCGETS, uintptr(unsafe.Pointer(&termios))); err != nil {
		return fmt.Errorf("ioctl(tty, tcgets): %s", err.Error())
	}

	termios.Oflag |= syscall.ONLCR

	if err := ioctl(terminal.Fd(), syscall.TCSETS, uintptr(unsafe.Pointer(&termios))); err != nil {
		return fmt.Errorf("ioctl(tty, tcsets): %s", err.Error())
	}

	return nil
}

func ioctl(fd uintptr, flag, data uintptr) error {
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, flag, data); err != 0 {
		return err
	}
	return nil
}
<commit_msg>Dadoo always closes pipes when it exits.<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/guardian/rundmc/dadoo"

	"github.com/eapache/go-resiliency/retrier"
	"github.com/opencontainers/runc/libcontainer/system"

	cmsg "github.com/opencontainers/runc/libcontainer/utils"
)

const MaxSocketDirPathLength = 80

func main() {
	os.Exit(run())
}

func run() int {
	tty := flag.Bool("tty", false, "tty requested")
	socketDirPath := flag.String("socket-dir-path", "", "path to a dir in which to store console sockets")
	runcRoot := flag.String("runc-root", "", "root directory for storage of container state (this should be located in tmpfs)")
	flag.Parse()

	runtime := flag.Args()[1] // e.g. runc
	processStateDir := flag.Args()[2]
	containerId := flag.Args()[3]

	signals := make(chan os.Signal, 100)
	signal.Notify(signals, syscall.SIGCHLD)

	runcExitCodePipe := os.NewFile(3, "/proc/self/fd/3")
	logFile := fmt.Sprintf("/proc/%d/fd/4", os.Getpid())
	logFD := os.NewFile(4, "/proc/self/fd/4")
	syncPipe := os.NewFile(5, "/proc/self/fd/5")
	pidFilePath := filepath.Join(processStateDir, "pidfile")

	stdinR, stdoutW, stderrW, winsz := openPipes(processStateDir)
	defer func() {
		tryClose(stdinR, stdoutW, stderrW, winsz)
	}()

	syncPipe.Write([]byte{0})

	stdoutR, stderrR := openStdioKeepAlivePipes(processStateDir)
	defer func() {
		tryClose(stdoutR, stderrR)
	}()

	ioWg := &sync.WaitGroup{}
	var runcExecCmd *exec.Cmd
	runtimeArgs := []string{"-debug", "-log", logFile}
	if *runcRoot != "" {
		runtimeArgs = append(runtimeArgs, "-root", *runcRoot)
	}
	runtimeArgs = append(runtimeArgs, "exec", "-d", "-p", fmt.Sprintf("/proc/%d/fd/0", os.Getpid()), "-pid-file", pidFilePath)
	if *tty {
		if len(*socketDirPath) > MaxSocketDirPathLength {
			return logAndExit(fmt.Sprintf("value for --socket-dir-path cannot exceed %d characters in length", MaxSocketDirPathLength))
		}
		ttySocketPath := setupTTYSocket(stdinR, stdoutW, winsz, pidFilePath, *socketDirPath, ioWg)
		runtimeArgs = append(runtimeArgs, "-tty", "-console-socket", ttySocketPath, containerId)
		runcExecCmd = exec.Command(runtime, runtimeArgs...)
	} else {
		runtimeArgs = append(runtimeArgs, containerId)
		runcExecCmd = exec.Command(runtime, runtimeArgs...)
		runcExecCmd.Stdin = stdinR
		runcExecCmd.Stdout = stdoutW
		runcExecCmd.Stderr = stderrW
	}

	// we need to be the subreaper so we can wait on the detached container process
	system.SetSubreaper(os.Getpid())

	if err := runcExecCmd.Start(); err != nil {
		runcExitCodePipe.Write([]byte{2})
		return 2
	}

	var status syscall.WaitStatus
	var rusage syscall.Rusage
	_, err := syscall.Wait4(runcExecCmd.Process.Pid, &status, 0, &rusage)
	check(err)    // Start succeeded but Wait4 failed, this can only be a programmer error
	logFD.Close() // No more logs from runc so close fd

	// also check that masterFD is received and streaming or whatevs
	runcExitCodePipe.Write([]byte{byte(status.ExitStatus())})
	if status.ExitStatus() != 0 {
		return 3 // nothing to wait for, container didn't launch
	}

	containerPid, err := parsePid(pidFilePath)
	check(err)

	return waitForContainerToExit(processStateDir, containerPid, signals, ioWg)
}

// If gdn server process dies, we need dadoo to keep stdout/err reader
// FDs so that Linux does not SIGPIPE the user process if it tries to use its end of
// these pipes.
func openStdioKeepAlivePipes(processStateDir string) (io.ReadCloser, io.ReadCloser) {
	keepStdoutAlive := openFifo(filepath.Join(processStateDir, "stdout"), os.O_RDONLY)
	keepStderrAlive := openFifo(filepath.Join(processStateDir, "stderr"), os.O_RDONLY)
	return keepStdoutAlive, keepStderrAlive
}

func waitForContainerToExit(processStateDir string, containerPid int, signals chan os.Signal, ioWg *sync.WaitGroup) (exitCode int) {
	for range signals {
		for {
			var status syscall.WaitStatus
			var rusage syscall.Rusage
			wpid, err := syscall.Wait4(-1, &status, syscall.WNOHANG, &rusage)
			if err != nil || wpid <= 0 {
				break // wait for next SIGCHLD
			}

			if wpid == containerPid {
				exitCode = status.ExitStatus()
				if status.Signaled() {
					exitCode = 128 + int(status.Signal())
				}

				ioWg.Wait() // wait for full output to be collected

				check(ioutil.WriteFile(filepath.Join(processStateDir, "exitcode"), []byte(strconv.Itoa(exitCode)), 0600))
				return exitCode
			}
		}
	}

	return logAndExit("ran out of signals") // cant happen
}

func openPipes(processStateDir string) (io.ReadCloser, io.WriteCloser, io.WriteCloser, io.ReadWriteCloser) {
	stdin := openFifo(filepath.Join(processStateDir, "stdin"), os.O_RDONLY)
	stdout := openFifo(filepath.Join(processStateDir, "stdout"), os.O_WRONLY|os.O_APPEND)
	stderr := openFifo(filepath.Join(processStateDir, "stderr"), os.O_WRONLY|os.O_APPEND)
	winsz := openFifo(filepath.Join(processStateDir, "winsz"), os.O_RDWR)
	openFifo(filepath.Join(processStateDir, "exit"), os.O_RDWR) // open just so guardian can detect it being closed when we exit

	return stdin, stdout, stderr, winsz
}

func openFifo(path string, flags int) io.ReadWriteCloser {
	r, err := os.OpenFile(path, flags, 0600)
	if os.IsNotExist(err) {
		return nil
	}

	check(err)
	return r
}

func setupTTYSocket(stdin io.Reader, stdout io.Writer, winszFifo io.Reader, pidFilePath, sockDirBase string, ioWg *sync.WaitGroup) string {
	sockDir, err := ioutil.TempDir(sockDirBase, "")
	check(err)

	ttySockPath := filepath.Join(sockDir, "tty.sock")
	l, err := net.Listen("unix", ttySockPath)
	check(err)

	//go to the background and set master
	go func(ln net.Listener) (err error) {
		// if any of the following errors, it means runc has connected to the
		// socket, so it must've started, thus we might need to kill the process
		defer func() {
			if err != nil {
				killProcess(pidFilePath)
				check(err)
			}
		}()

		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer conn.Close()

		// Close ln, to allow for other instances to take over.
		ln.Close()

		// Get the fd of the connection.
		unixconn, ok := conn.(*net.UnixConn)
		if !ok {
			return
		}

		socket, err := unixconn.File()
		if err != nil {
			return
		}
		defer socket.Close()

		// Get the master file descriptor from runC.
		master, err := cmsg.RecvFd(socket)
		if err != nil {
			return
		}

		os.RemoveAll(sockDir)
		if err = setOnlcr(master); err != nil {
			return
		}
		streamProcess(master, stdin, stdout, winszFifo, ioWg)

		return
	}(l)

	return ttySockPath
}

func streamProcess(m *os.File, stdin io.Reader, stdout io.Writer, winszFifo io.Reader, ioWg *sync.WaitGroup) {
	ioWg.Add(1)
	go func() {
		defer ioWg.Done()
		io.Copy(stdout, m)
	}()

	go io.Copy(m, stdin)

	go func() {
		for {
			var winSize garden.WindowSize
			if err := json.NewDecoder(winszFifo).Decode(&winSize); err != nil {
				fmt.Printf("invalid winsz event: %s\n", err)
				continue // not much we can do here..
			}
			dadoo.SetWinSize(m, winSize)
		}
	}()
}

func killProcess(pidFilePath string) {
	pid, err := readPid(pidFilePath)
	if err == nil {
		syscall.Kill(pid, syscall.SIGKILL)
	}
}

func readPid(pidFilePath string) (int, error) {
	retrier := retrier.New(retrier.ConstantBackoff(20, 500*time.Millisecond), nil)
	var (
		pid int = -1
		err error
	)
	retrier.Run(func() error {
		pid, err = parsePid(pidFilePath)
		return err
	})

	return pid, err
}

func parsePid(pidFile string) (int, error) {
	b, err := ioutil.ReadFile(pidFile)
	if err != nil {
		return -1, err
	}

	var pid int
	if _, err := fmt.Sscanf(string(b), "%d", &pid); err != nil {
		return -1, err
	}

	return pid, nil
}

func logAndExit(msg string) int {
	fmt.Println(msg)
	return 2
}

func check(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(2)
	}
}

func tryClose(closers ...io.Closer) {
	for _, closer := range closers {
		if closer != nil {
			closer.Close()
		}
	}
}

// setOnlcr copied from runc
// https://github.com/cloudfoundry-incubator/runc/blob/02ec89829b24dfce45bb207d2344e0e6d078a93c/libcontainer/console_linux.go#L144-L160
func setOnlcr(terminal *os.File) error {
	var termios syscall.Termios

	if err := ioctl(terminal.Fd(), syscall.TCGETS, uintptr(unsafe.Pointer(&termios))); err != nil {
		return fmt.Errorf("ioctl(tty, tcgets): %s", err.Error())
	}

	termios.Oflag |= syscall.ONLCR

	if err := ioctl(terminal.Fd(), syscall.TCSETS, uintptr(unsafe.Pointer(&termios))); err != nil {
		return fmt.Errorf("ioctl(tty, tcsets): %s", err.Error())
	}

	return nil
}

func ioctl(fd uintptr, flag, data uintptr) error {
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, flag, data); err != 0 {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 ePoxy Authors
//
// 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.
//////////////////////////////////////////////////////////////////////////////

// A minimal client for adding Host records to Datastore for testing. This
// command is ONLY for testing. Host record management by direct access to
// Datastore will not be supported by ePoxy.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"

	"cloud.google.com/go/datastore"
	"github.com/kr/pretty"
	"github.com/m-lab/epoxy/storage"
)

const usage = `USAGE:
**Only use for testing.**

EXAMPLE:
    epoxy_admin --project mlab-sandbox \
        --hostname mlab3.iad1t.measurement-lab.org \
        --address 165.117.240.35 \
        --stage1 https://storage.googleapis.com/epoxy-mlab-sandbox/os/stage1to2.ipxe
        --stage2 https://storage.googleapis.com/epoxy-mlab-sandbox/os/stage2to3.json
        --stage3 https://storage.googleapis.com/epoxy-mlab-sandbox/os/stage3post.json
`

var (
	fProject  string
	fHostname string
	fAddress  string
	fStage1   string
	fStage2   string
	fStage3   string
)

func init() {
	// Add an alternate usage message.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, usage)
		flag.PrintDefaults()
	}
	flag.StringVar(&fProject, "project", "mlab-sandbox", "GCP project ID.")
	flag.StringVar(&fHostname, "hostname", "mlab3.iad1t.measurement-lab.org", "Hostname of new record.")
	flag.StringVar(&fAddress, "address", "165.117.240.35", "IP address of hostname.")
	flag.StringVar(&fStage1, "stage1",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_coreos/stage1to2.ipxe",
		"Absolute URL to an action definition to run during stage1 to boot stage2.")
	flag.StringVar(&fStage2, "stage2",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_coreos/stage2to3.json",
		"Absolute URL to an action definition to run during stage2 to boot stage3.")
	flag.StringVar(&fStage3, "stage3",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_coreos/stage3post.json",
		"Absolute URL to an action definition to run after booting stage3.")
}

func main() {
	flag.Parse()

	// Setup Datastore client.
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, fProject)
	if err != nil {
		log.Fatalf("Failed to create new datastore client: %s", err)
	}

	// Save the host record to Datstore.
	ds := storage.NewDatastoreConfig(client)
	h := &storage.Host{
		Name:     fHostname,
		IPv4Addr: fAddress,
		Boot: storage.Sequence{
			Stage1ChainURL: fStage1,
			Stage2ChainURL: fStage2,
			Stage3ChainURL: fStage3,
		},
	}
	if err = ds.Save(h); err != nil {
		log.Fatalf("%s", err)
	}

	// Retrieve the host record from Datastore to exercise the full save & load path.
	h2, err := ds.Load(h.Name)
	if err != nil {
		log.Fatalf("%s", err)
	}
	pretty.Print(h2.String())
}
<commit_msg>Update epoxy admin to support update operations<commit_after>// Copyright 2016 ePoxy Authors
//
// 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.
//////////////////////////////////////////////////////////////////////////////

// A minimal client for adding Host records to Datastore for testing. This
// command is ONLY for testing. Host record management by direct access to
// Datastore will not be supported by ePoxy.
//
// TODO:
//   * Create distinct subcommands, e.g. create, update, delete.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"

	"cloud.google.com/go/datastore"
	"github.com/kr/pretty"
	"github.com/m-lab/epoxy/storage"
)

const usage = `USAGE:
**ONLY USE FOR TESTING**

EXAMPLE:
    epoxy_admin --project mlab-sandbox \
        --hostname mlab3.iad1t.measurement-lab.org \
        --address 165.117.240.35 \
        --boot-stage1 https://storage.googleapis.com/epoxy-mlab-sandbox/os/stage1to2.ipxe
        --boot-stage2 https://storage.googleapis.com/epoxy-mlab-sandbox/os/stage2to3.json
        --boot-stage3 https://storage.googleapis.com/epoxy-mlab-sandbox/os/stage3post.json
`

var (
	fProject      string
	fHostname     string
	fAddress      string
	fUpdate       bool
	fBootStage1   string
	fBootStage2   string
	fBootStage3   string
	fUpdateStage1 string
	fUpdateStage2 string
	fUpdateStage3 string
)

func init() {
	// Add an alternate usage message.
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, usage)
		flag.PrintDefaults()
	}
	flag.StringVar(&fProject, "project", "mlab-sandbox", "GCP project ID.")
	flag.StringVar(&fHostname, "hostname", "mlab3.iad1t.measurement-lab.org", "Hostname of new record.")
	flag.StringVar(&fAddress, "address", "165.117.240.35", "IP address of hostname.")
	flag.BoolVar(&fUpdate, "update", false,
		"Set Host.UpdateEnabled to true for an existing Host. Do not specify when creating a new Host.")
	flag.StringVar(&fBootStage1, "boot-stage1",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_coreos/stage1to2.ipxe",
		"Absolute URL to an action definition to run during stage1 to boot stage2.")
	flag.StringVar(&fBootStage2, "boot-stage2",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_coreos/stage2to3.json",
		"Absolute URL to an action definition to run during stage2 to boot stage3.")
	flag.StringVar(&fBootStage3, "boot-stage3",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_coreos/stage3post.json",
		"Absolute URL to an action definition to run after booting stage3.")
	flag.StringVar(&fUpdateStage1, "update-stage1",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_mlxupdate/stage1to2.ipxe",
		"Absolute URL to an action definition to run during stage1 to boot stage2.")
	flag.StringVar(&fUpdateStage2, "update-stage2",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_mlxupdate/stage2to3.json",
		"Absolute URL to an action definition to run during stage2 to boot stage3.")
	flag.StringVar(&fUpdateStage3, "update-stage3",
		"https://storage.googleapis.com/epoxy-mlab-sandbox/stage3_mlxupdate/stage3post.json",
		"Absolute URL to an action definition to run after booting stage3.")
}

func main() {
	flag.Parse()

	// Setup Datastore client.
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, fProject)
	if err != nil {
		log.Fatalf("Failed to create new datastore client: %s", err)
	}

	// Save the host record to Datstore.
	ds := storage.NewDatastoreConfig(client)
	var h *storage.Host
	var err error

	if fUpdate {
		// Retrieve the host record from Datastore before updating it.
		h, err = ds.Load(fHostname)
		if err != nil {
			log.Fatalf("%s", err)
		}
		h.UpdateEnabled = true

	} else {

		// Create a new host record.
		h = &storage.Host{
			Name:          fHostname,
			IPv4Addr:      fAddress,
			UpdateEnabled: fUpdate,
			Boot: storage.Sequence{
				Stage1ChainURL: fBootStage1,
				Stage2ChainURL: fBootStage2,
				Stage3ChainURL: fBootStage3,
			},
			Update: storage.Sequence{
				Stage1ChainURL: fUpdateStage1,
				Stage2ChainURL: fUpdateStage2,
				Stage3ChainURL: fUpdateStage3,
			},
		}
	}
	if err = ds.Save(h); err != nil {
		log.Fatalf("%s", err)
	}

	// Retrieve the host record from Datastore to exercise the full save & load path.
	h2, err := ds.Load(h.Name)
	if err != nil {
		log.Fatalf("%s", err)
	}
	pretty.Print(h2.String())
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/gofrs/flock"
)

// version is the program's version number.
var version = "unknown"

func main() {
	log.SetFlags(0)

	args := getArgs()

	config, err := NewConfig(args.ConfigFile, args.DatabaseDirectory)
	if err != nil {
		fatal(args, "Error loading configuration file", err)
	}
	if args.Verbose {
		log.Printf("Using config file %s", args.ConfigFile)
		log.Printf("Using database directory %s", config.DatabaseDirectory)
	}

	lock, err := setup(config, args.Verbose)
	if err != nil {
		fatal(args, "Error preparing to update", err)
	}
	defer func() {
		if err := lock.Unlock(); err != nil {
			fatal(args, "Error unlocking lock file", errors.Wrap(err, "unlocking"))
		}
	}()

	if err := run(config, args.Verbose); err != nil {
		fatal(args, "Error retrieving updates", err)
	}
}

func fatal(
	args *Args,
	msg string,
	err error,
) {
	if args.StackTrace {
		log.Print(msg + fmt.Sprintf(": %+v", err))
	} else {
		log.Print(msg + fmt.Sprintf(": %s", err))
	}
	os.Exit(1)
}

func setup(
	config *Config,
	verbose bool,
) (*flock.Flock, error) {
	if err := maybeSetProxy(config, verbose); err != nil {
		return nil, err
	}

	if err := checkEnvironment(config); err != nil {
		return nil, err
	}

	lock := flock.New(config.LockFile)
	ok, err := lock.TryLock()
	if err != nil {
		return nil, errors.Wrap(err, "error acquiring a lock")
	}
	if !ok {
		return nil, errors.Errorf("could not acquire lock on %s", config.LockFile)
	}
	if verbose {
		log.Printf("Acquired lock file lock (%s)", config.LockFile)
	}

	return lock, nil
}

// Do not set a timeout to allow for very slow connections. Note the client
// will have TCP KeepAlive's enabled by default due to using
// http.DefaultTransport (which uses a net.Dialer with KeepAlive set).
var client = &http.Client{}

func maybeSetProxy(
	config *Config,
	verbose bool,
) error {
	if config.Proxy == nil {
		return nil
	}

	if verbose {
		log.Printf("Using proxy: %s", config.Proxy)
	}
	http.DefaultTransport.(*http.Transport).Proxy = http.ProxyURL(config.Proxy)

	return nil
}

func checkEnvironment(
	config *Config,
) error {
	fi, err := os.Stat(config.DatabaseDirectory)
	if err != nil {
		return errors.Wrap(err, "database directory is not available")
	}

	if !fi.IsDir() {
		return errors.New("database directory is not a directory")
	}

	// I don't think there is a reliable cross platform way to check the
	// directory is writable. We'll discover that when we try to write to it
	// anyway.

	return nil
}

func run(
	config *Config,
	verbose bool,
) error {
	for _, editionID := range config.EditionIDs {
		if err := updateEdition(config, verbose, editionID); err != nil {
			return errors.WithMessage(err, "error updating "+editionID)
		}
	}
	return nil
}

func updateEdition(
	config *Config,
	verbose bool,
	editionID string,
) error {
	filename, err := getFilename(config, verbose, editionID)
	if err != nil {
		return errors.WithMessage(err, "error retrieving filename")
	}

	md5, err := getCurrentMD5(config, verbose, filename)
	if err != nil {
		return errors.WithMessage(err, "error retrieving current MD5 of "+filename)
	}

	if err := maybeUpdate(
		config,
		verbose,
		editionID,
		filename,
		md5,
	); err != nil {
		return errors.WithMessage(err, "error updating")
	}

	return nil
}

func getFilename(
	config *Config,
	verbose bool,
	editionID string,
) (string, error) {
	url := fmt.Sprintf(
		"%s/app/update_getfilename?product_id=%s",
		config.URL,
		url.QueryEscape(editionID),
	)

	if verbose {
		log.Printf("Performing get filename request to %s", url)
	}
	res, err := client.Get(url)
	if err != nil {
		return "", errors.Wrap(err, "error performing HTTP request")
	}
	defer func() {
		if err := res.Body.Close(); err != nil {
			log.Fatalf("Error closing response body: %+v", errors.Wrap(err, "closing body"))
		}
	}()

	buf, err := ioutil.ReadAll(io.LimitReader(res.Body, 256))
	if err != nil {
		return "", errors.Wrap(err, "error reading response body")
	}

	if res.StatusCode != http.StatusOK {
		return "", errors.Errorf("unexpected HTTP status code: %s: %s", res.Status, buf)
	}

	if len(buf) == 0 {
		return "", errors.New("response body is empty")
	}

	if bytes.Count(buf, []byte("\n")) > 0 ||
		bytes.Count(buf, []byte("\x00")) > 0 {
		return "", errors.New("invalid characters in filename")
	}

	return string(buf), nil
}

const zeroMD5 = "00000000000000000000000000000000"

func getCurrentMD5(
	config *Config,
	verbose bool,
	filename string,
) (string, error) {
	path := filepath.Join(config.DatabaseDirectory, filename)

	fh, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			if verbose {
				log.Printf("Not calculating MD5 sum as file does not exist: %s", path)
			}
			return zeroMD5, nil
		}
		return "", errors.Wrap(err, "error opening file")
	}
	defer func() {
		if err := fh.Close(); err != nil {
			log.Fatalf("Error closing file: %+v", errors.Wrap(err, "closing file"))
		}
	}()

	fi, err := fh.Stat()
	if err != nil {
		return "", errors.Wrap(err, "error stat'ing file")
	}
	if !fi.Mode().IsRegular() {
		return "", errors.New("not a regular file")
	}

	h := md5.New()
	if _, err := io.Copy(h, fh); err != nil {
		return "", errors.Wrap(err, "error reading file")
	}
	sum := fmt.Sprintf("%x", h.Sum(nil))
	if verbose {
		log.Printf("Calculated MD5 sum for %s: %s", path, sum)
	}
	return sum, nil
}

func maybeUpdate(
	config *Config,
	verbose bool,
	editionID,
	filename,
	md5 string,
) error {
	url := fmt.Sprintf(
		"%s/geoip/databases/%s/update?db_md5=%s",
		config.URL,
		url.PathEscape(editionID),
		url.QueryEscape(md5),
	)

	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return errors.Wrap(err, "error creating request")
	}
	if config.AccountID != 0 {
		req.SetBasicAuth(fmt.Sprintf("%d", config.AccountID), config.LicenseKey)
	}

	if verbose {
		log.Printf("Performing update request to %s", url)
	}
	res, err := client.Do(req)
	if err != nil {
		return errors.Wrap(err, "error performing HTTP request")
	}
	defer func() {
		if err := res.Body.Close(); err != nil {
			log.Fatalf("Error closing response body: %+v", errors.Wrap(err, "closing body"))
		}
	}()

	if res.StatusCode == http.StatusNotModified {
		if verbose {
			log.Printf("No new updates available for %s", editionID)
		}
		return nil
	}

	if res.StatusCode != http.StatusOK {
		buf, err := ioutil.ReadAll(io.LimitReader(res.Body, 256))
		if err == nil {
			return errors.Errorf("unexpected HTTP status code: %s: %s", res.Status, buf)
		}
		return errors.Errorf("unexpected HTTP status code: %s", res.Status)
	}

	newMD5 := res.Header.Get("X-Database-MD5")
	if newMD5 == "" {
		return errors.New("no X-Database-MD5 header found")
	}
	lastModified, err := getLastModified(res.Header)
	if err != nil {
		return err
	}

	return writeAndCheck(config, verbose, filename, res.Body, newMD5, lastModified)
}

func getLastModified(
	headers http.Header,
) (time.Time, error) {
	lastModifiedStr := headers.Get("Last-Modified")
	if lastModifiedStr == "" {
		return time.Time{}, errors.New("no Last-Modified header found")
	}

	t, err := time.ParseInLocation(time.RFC1123, lastModifiedStr, time.UTC)
	if err != nil {
		return time.Time{}, errors.Wrap(err, "error parsing time")
	}

	return t, nil
}

func writeAndCheck(
	config *Config,
	verbose bool,
	filename string,
	body io.Reader,
	newMD5 string,
	lastModified time.Time,
) error {
	targetTest := filepath.Join(
		config.DatabaseDirectory,
		fmt.Sprintf("%s.test", filename),
	)

	fh, err := os.OpenFile(targetTest, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return errors.Wrap(err, "error creating file")
	}

	gzReader, err := gzip.NewReader(body)
	if err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error creating gzip reader")
	}

	md5Writer := md5.New()
	multiWriter := io.MultiWriter(fh, md5Writer)

	if _, err := io.Copy(multiWriter, gzReader); err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		_ = gzReader.Close()
		return errors.Wrap(err, "error reading/writing")
	}

	if err := gzReader.Close(); err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error closing gzip reader")
	}

	if err := fh.Sync(); err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error syncing file")
	}

	if err := fh.Close(); err != nil {
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error closing file")
	}

	gotMD5 := fmt.Sprintf("%x", md5Writer.Sum(nil))
	if !strings.EqualFold(gotMD5, newMD5) {
		_ = os.Remove(targetTest)
		return errors.Errorf("MD5 of new database (%s) does not match expected MD5 (%s)",
			gotMD5, newMD5)
	}

	target := filepath.Join(config.DatabaseDirectory, filename)

	if err := os.Rename(targetTest, target); err != nil {
		_ = os.Remove(targetTest)
		return errors.New("error moving database into place")
	}

	if config.PreserveFileTimes {
		if err := os.Chtimes(target, lastModified, lastModified); err != nil {
			return errors.Wrap(err, "error setting times on file")
		}
	}

	// fsync the directory. http://austingroupbugs.net/view.php?id=672

	dh, err := os.Open(config.DatabaseDirectory)
	if err != nil {
		return errors.Wrap(err, "error opening database directory")
	}
	defer func() {
		if err := dh.Close(); err != nil {
			log.Fatalf("Error closing directory: %+v", errors.Wrap(err, "closing directory"))
		}
	}()

	if err := dh.Sync(); err != nil {
		return errors.Wrap(err, "error syncing database directory")
	}

	if verbose {
		log.Printf("Updated %s", target)
	}
	return nil
}
<commit_msg>Set version from build info, if available<commit_after>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime/debug"
	"strings"
	"time"

	"github.com/gofrs/flock"
	"github.com/pkg/errors"
)

// version is the program's version number.
var version = "unknown"

func main() {
	if info, ok := debug.ReadBuildInfo(); ok && info.Main.Version != "(devel)" {
		version = info.Main.Version
	}

	log.SetFlags(0)

	args := getArgs()

	config, err := NewConfig(args.ConfigFile, args.DatabaseDirectory)
	if err != nil {
		fatal(args, "Error loading configuration file", err)
	}
	if args.Verbose {
		log.Printf("Using config file %s", args.ConfigFile)
		log.Printf("Using database directory %s", config.DatabaseDirectory)
	}

	lock, err := setup(config, args.Verbose)
	if err != nil {
		fatal(args, "Error preparing to update", err)
	}
	defer func() {
		if err := lock.Unlock(); err != nil {
			fatal(args, "Error unlocking lock file", errors.Wrap(err, "unlocking"))
		}
	}()

	if err := run(config, args.Verbose); err != nil {
		fatal(args, "Error retrieving updates", err)
	}
}

func fatal(
	args *Args,
	msg string,
	err error,
) {
	if args.StackTrace {
		log.Print(msg + fmt.Sprintf(": %+v", err))
	} else {
		log.Print(msg + fmt.Sprintf(": %s", err))
	}
	os.Exit(1)
}

func setup(
	config *Config,
	verbose bool,
) (*flock.Flock, error) {
	if err := maybeSetProxy(config, verbose); err != nil {
		return nil, err
	}

	if err := checkEnvironment(config); err != nil {
		return nil, err
	}

	lock := flock.New(config.LockFile)
	ok, err := lock.TryLock()
	if err != nil {
		return nil, errors.Wrap(err, "error acquiring a lock")
	}
	if !ok {
		return nil, errors.Errorf("could not acquire lock on %s", config.LockFile)
	}
	if verbose {
		log.Printf("Acquired lock file lock (%s)", config.LockFile)
	}

	return lock, nil
}

// Do not set a timeout to allow for very slow connections. Note the client
// will have TCP KeepAlive's enabled by default due to using
// http.DefaultTransport (which uses a net.Dialer with KeepAlive set).
var client = &http.Client{}

func maybeSetProxy(
	config *Config,
	verbose bool,
) error {
	if config.Proxy == nil {
		return nil
	}

	if verbose {
		log.Printf("Using proxy: %s", config.Proxy)
	}
	http.DefaultTransport.(*http.Transport).Proxy = http.ProxyURL(config.Proxy)

	return nil
}

func checkEnvironment(
	config *Config,
) error {
	fi, err := os.Stat(config.DatabaseDirectory)
	if err != nil {
		return errors.Wrap(err, "database directory is not available")
	}

	if !fi.IsDir() {
		return errors.New("database directory is not a directory")
	}

	// I don't think there is a reliable cross platform way to check the
	// directory is writable. We'll discover that when we try to write to it
	// anyway.

	return nil
}

func run(
	config *Config,
	verbose bool,
) error {
	for _, editionID := range config.EditionIDs {
		if err := updateEdition(config, verbose, editionID); err != nil {
			return errors.WithMessage(err, "error updating "+editionID)
		}
	}
	return nil
}

func updateEdition(
	config *Config,
	verbose bool,
	editionID string,
) error {
	filename, err := getFilename(config, verbose, editionID)
	if err != nil {
		return errors.WithMessage(err, "error retrieving filename")
	}

	md5, err := getCurrentMD5(config, verbose, filename)
	if err != nil {
		return errors.WithMessage(err, "error retrieving current MD5 of "+filename)
	}

	if err := maybeUpdate(
		config,
		verbose,
		editionID,
		filename,
		md5,
	); err != nil {
		return errors.WithMessage(err, "error updating")
	}

	return nil
}

func getFilename(
	config *Config,
	verbose bool,
	editionID string,
) (string, error) {
	url := fmt.Sprintf(
		"%s/app/update_getfilename?product_id=%s",
		config.URL,
		url.QueryEscape(editionID),
	)

	if verbose {
		log.Printf("Performing get filename request to %s", url)
	}
	res, err := client.Get(url)
	if err != nil {
		return "", errors.Wrap(err, "error performing HTTP request")
	}
	defer func() {
		if err := res.Body.Close(); err != nil {
			log.Fatalf("Error closing response body: %+v", errors.Wrap(err, "closing body"))
		}
	}()

	buf, err := ioutil.ReadAll(io.LimitReader(res.Body, 256))
	if err != nil {
		return "", errors.Wrap(err, "error reading response body")
	}

	if res.StatusCode != http.StatusOK {
		return "", errors.Errorf("unexpected HTTP status code: %s: %s", res.Status, buf)
	}

	if len(buf) == 0 {
		return "", errors.New("response body is empty")
	}

	if bytes.Count(buf, []byte("\n")) > 0 ||
		bytes.Count(buf, []byte("\x00")) > 0 {
		return "", errors.New("invalid characters in filename")
	}

	return string(buf), nil
}

const zeroMD5 = "00000000000000000000000000000000"

func getCurrentMD5(
	config *Config,
	verbose bool,
	filename string,
) (string, error) {
	path := filepath.Join(config.DatabaseDirectory, filename)

	fh, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			if verbose {
				log.Printf("Not calculating MD5 sum as file does not exist: %s", path)
			}
			return zeroMD5, nil
		}
		return "", errors.Wrap(err, "error opening file")
	}
	defer func() {
		if err := fh.Close(); err != nil {
			log.Fatalf("Error closing file: %+v", errors.Wrap(err, "closing file"))
		}
	}()

	fi, err := fh.Stat()
	if err != nil {
		return "", errors.Wrap(err, "error stat'ing file")
	}
	if !fi.Mode().IsRegular() {
		return "", errors.New("not a regular file")
	}

	h := md5.New()
	if _, err := io.Copy(h, fh); err != nil {
		return "", errors.Wrap(err, "error reading file")
	}
	sum := fmt.Sprintf("%x", h.Sum(nil))
	if verbose {
		log.Printf("Calculated MD5 sum for %s: %s", path, sum)
	}
	return sum, nil
}

func maybeUpdate(
	config *Config,
	verbose bool,
	editionID,
	filename,
	md5 string,
) error {
	url := fmt.Sprintf(
		"%s/geoip/databases/%s/update?db_md5=%s",
		config.URL,
		url.PathEscape(editionID),
		url.QueryEscape(md5),
	)

	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return errors.Wrap(err, "error creating request")
	}
	if config.AccountID != 0 {
		req.SetBasicAuth(fmt.Sprintf("%d", config.AccountID), config.LicenseKey)
	}

	if verbose {
		log.Printf("Performing update request to %s", url)
	}
	res, err := client.Do(req)
	if err != nil {
		return errors.Wrap(err, "error performing HTTP request")
	}
	defer func() {
		if err := res.Body.Close(); err != nil {
			log.Fatalf("Error closing response body: %+v", errors.Wrap(err, "closing body"))
		}
	}()

	if res.StatusCode == http.StatusNotModified {
		if verbose {
			log.Printf("No new updates available for %s", editionID)
		}
		return nil
	}

	if res.StatusCode != http.StatusOK {
		buf, err := ioutil.ReadAll(io.LimitReader(res.Body, 256))
		if err == nil {
			return errors.Errorf("unexpected HTTP status code: %s: %s", res.Status, buf)
		}
		return errors.Errorf("unexpected HTTP status code: %s", res.Status)
	}

	newMD5 := res.Header.Get("X-Database-MD5")
	if newMD5 == "" {
		return errors.New("no X-Database-MD5 header found")
	}
	lastModified, err := getLastModified(res.Header)
	if err != nil {
		return err
	}

	return writeAndCheck(config, verbose, filename, res.Body, newMD5, lastModified)
}

func getLastModified(
	headers http.Header,
) (time.Time, error) {
	lastModifiedStr := headers.Get("Last-Modified")
	if lastModifiedStr == "" {
		return time.Time{}, errors.New("no Last-Modified header found")
	}

	t, err := time.ParseInLocation(time.RFC1123, lastModifiedStr, time.UTC)
	if err != nil {
		return time.Time{}, errors.Wrap(err, "error parsing time")
	}

	return t, nil
}

func writeAndCheck(
	config *Config,
	verbose bool,
	filename string,
	body io.Reader,
	newMD5 string,
	lastModified time.Time,
) error {
	targetTest := filepath.Join(
		config.DatabaseDirectory,
		fmt.Sprintf("%s.test", filename),
	)

	fh, err := os.OpenFile(targetTest, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return errors.Wrap(err, "error creating file")
	}

	gzReader, err := gzip.NewReader(body)
	if err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error creating gzip reader")
	}

	md5Writer := md5.New()
	multiWriter := io.MultiWriter(fh, md5Writer)

	if _, err := io.Copy(multiWriter, gzReader); err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		_ = gzReader.Close()
		return errors.Wrap(err, "error reading/writing")
	}

	if err := gzReader.Close(); err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error closing gzip reader")
	}

	if err := fh.Sync(); err != nil {
		_ = fh.Close()
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error syncing file")
	}

	if err := fh.Close(); err != nil {
		_ = os.Remove(targetTest)
		return errors.Wrap(err, "error closing file")
	}

	gotMD5 := fmt.Sprintf("%x", md5Writer.Sum(nil))
	if !strings.EqualFold(gotMD5, newMD5) {
		_ = os.Remove(targetTest)
		return errors.Errorf("MD5 of new database (%s) does not match expected MD5 (%s)",
			gotMD5, newMD5)
	}

	target := filepath.Join(config.DatabaseDirectory, filename)

	if err := os.Rename(targetTest, target); err != nil {
		_ = os.Remove(targetTest)
		return errors.New("error moving database into place")
	}

	if config.PreserveFileTimes {
		if err := os.Chtimes(target, lastModified, lastModified); err != nil {
			return errors.Wrap(err, "error setting times on file")
		}
	}

	// fsync the directory. http://austingroupbugs.net/view.php?id=672

	dh, err := os.Open(config.DatabaseDirectory)
	if err != nil {
		return errors.Wrap(err, "error opening database directory")
	}
	defer func() {
		if err := dh.Close(); err != nil {
			log.Fatalf("Error closing directory: %+v", errors.Wrap(err, "closing directory"))
		}
	}()

	if err := dh.Sync(); err != nil {
		return errors.Wrap(err, "error syncing database directory")
	}

	if verbose {
		log.Printf("Updated %s", target)
	}
	return nil
}
<|endoftext|>