<commit_before>package encoders

import (
	"errors"
	// "fmt"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)

	// create input context
	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	// create output context
	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	//get audio and video stream and the streaMap
	streamMap, srcVideoStream, srcAudioStream, err := getAudioVideoStreams(inputCtx, outputCtx, job)
	if err != nil {
		return err
	}
	//calculate total number of frames
	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())
	//process all frames and update the job progress
	err = processAllFramesAndUpdateJobProgress(inputCtx, outputCtx, streamMap, job, dbInstance, totalFrames)
	if err != nil {
		return err
	}

	processNewFrames(inputCtx, outputCtx, streamMap)
	if err != nil {
		return err
	}

	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

func processNewFrames(inputCtx *gmf.FmtCtx, outputCtx *gmf.FmtCtx, streamMap map[int]int) error {
	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		_, ost, err := getOutputAndInputStream(inputCtx, outputCtx, streamMap, 0)
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(ost.CodecCtx()); ready {
				configurePacket(p, ost, frame)
				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}
			ost.Pts++
		}

		gmf.Release(frame)
	}

	return nil
}

func processAllFramesAndUpdateJobProgress(inputCtx *gmf.FmtCtx, outputCtx *gmf.FmtCtx, streamMap map[int]int, job types.Job, dbInstance db.Storage, totalFrames float64) error {
	var lastDelta int64
	for packet := range inputCtx.GetNewPackets() {
		ist, ost, err := getOutputAndInputStream(inputCtx, outputCtx, streamMap, packet.StreamIndex())
		if err != nil {
			return err
		}
		framesCount := float64(0)
		for frame := range packet.Frames(ist.CodecCtx()) {
			err := proccessFrame(ist, ost, packet, frame, outputCtx, &lastDelta)
			if err != nil {
				return err
			}

			ost.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}

		gmf.Release(packet)
	}
	return nil
}

func getOutputAndInputStream(inputCtx *gmf.FmtCtx, outputCtx *gmf.FmtCtx, streamMap map[int]int, inputIndex int) (*gmf.Stream, *gmf.Stream, error) {
	inputStream, err := inputCtx.GetStream(inputIndex)
	if err != nil {
		return nil, nil, err
	}
	outputStream, err := outputCtx.GetStream(streamMap[inputStream.Index()])
	if err != nil {
		return nil, nil, err
	}

	return inputStream, outputStream, nil
}

func getAudioVideoStreams(inputCtx *gmf.FmtCtx, outputCtx *gmf.FmtCtx, job types.Job) (map[int]int, *gmf.Stream, *gmf.Stream, error) {
	streamMap := make(map[int]int, 0)

	// add video stream to streamMap
	srcVideoStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)
	if err != nil {
		return nil, nil, nil, errors.New("unable to get the best video stream inside the input context")
	}
	videoCodec := getVideoCodec(job)
	inputIndex, outputIndex, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		return nil, nil, nil, err
	}
	streamMap[inputIndex] = outputIndex

	// add audio stream to streamMap
	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return nil, nil, nil, errors.New("unable to get the best audio stream inside the input context")
	}
	audioCodec := getAudioCodec(job)
	inputIndex, outputIndex, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return nil, nil, nil, err
	}
	streamMap[inputIndex] = outputIndex
	if err := outputCtx.WriteHeader(); err != nil {
		return nil, nil, nil, err
	}

	return streamMap, srcVideoStream, srcAudioStream, nil
}

func configureAudioFrame(packet *gmf.Packet, inputStream *gmf.Stream, outputStream *gmf.Stream, frame *gmf.Frame, lastDelta *int64) {
	fsTb := gmf.AVR{Num: 1, Den: inputStream.CodecCtx().SampleRate()}
	outTb := gmf.AVR{Num: 1, Den: inputStream.CodecCtx().SampleRate()}

	frame.SetPts(packet.Pts())

	pts := gmf.RescaleDelta(inputStream.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), lastDelta, outTb.AVRational())

	frame.SetNbSamples(outputStream.CodecCtx().FrameSize())
	frame.SetFormat(outputStream.CodecCtx().SampleFmt())
	frame.SetChannelLayout(outputStream.CodecCtx().ChannelLayout())
	frame.SetPts(pts)
}

func configurePacket(packet *gmf.Packet, outputStream *gmf.Stream, frame *gmf.Frame) *gmf.Packet {
	if packet.Pts() != gmf.AV_NOPTS_VALUE {
		packet.SetPts(gmf.RescaleQ(packet.Pts(), outputStream.CodecCtx().TimeBase(), outputStream.TimeBase()))
	}

	if packet.Dts() != gmf.AV_NOPTS_VALUE {
		packet.SetDts(gmf.RescaleQ(packet.Dts(), outputStream.CodecCtx().TimeBase(), outputStream.TimeBase()))
	}

	packet.SetStreamIndex(outputStream.Index())

	return packet
}

func proccessFrame(inputStream *gmf.Stream, outputStream *gmf.Stream, packet *gmf.Packet, frame *gmf.Frame, outputCtx *gmf.FmtCtx, lastDelta *int64) error {
	if outputStream.IsAudio() {
		configureAudioFrame(packet, inputStream, outputStream, frame, lastDelta)
	} else {
		frame.SetPts(outputStream.Pts)
	}

	if newPacket, ready, _ := frame.EncodeNewPacket(outputStream.CodecCtx()); ready {
		configurePacket(newPacket, outputStream, frame)
		if err := outputCtx.WritePacket(newPacket); err != nil {
			return err
		}
		gmf.Release(newPacket)
	}

	return nil
}

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, inputStream *gmf.Stream) (int, int, error) {
	var codecContext *gmf.CodecCtx
	var outputStream *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if outputStream = oc.NewStream(codec); outputStream == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(outputStream)

	if codecContext = gmf.NewCodecCtx(codec); codecContext == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(codecContext)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		codecContext.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		codecContext.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		err := setAudioCtxParams(codecContext, inputStream, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		err := setVideoCtxParams(codecContext, inputStream, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if err := codecContext.Open(nil); err != nil {
		return 0, 0, err
	}

	outputStream.SetCodecCtx(codecContext)

	return inputStream.Index(), outputStream.Index(), nil
}

func getProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile != "" {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

func getVideoCodec(job types.Job) string {
	codecs := map[string]string{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
		"aac":    "aac",
	}

	if codec, ok := codecs[job.Preset.Video.Codec]; ok {
		return codec
	}
	return "libx264"
}

func getAudioCodec(job types.Job) string {
	codecs := map[string]string{
		"aac":    "aac",
		"vorbis": "vorbis",
	}
	if codec, ok := codecs[job.Preset.Audio.Codec]; ok {
		return codec
	}
	return "aac"
}

func GetResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}

func setAudioCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetBitRate(bitrate)
	codecContext.SetSampleFmt(ist.CodecCtx().SampleFmt())
	codecContext.SetSampleRate(ist.CodecCtx().SampleRate())
	codecContext.SetChannels(ist.CodecCtx().Channels())
	codecContext.SelectChannelLayout()
	codecContext.SelectSampleRate()
	return nil
}

func setVideoCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	codecContext.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

	if job.Preset.Video.Codec == "h264" {
		profile := getProfile(job)
		codecContext.SetProfile(profile)
	}

	gop, err := strconv.Atoi(job.Preset.Video.GopSize)
	if err != nil {
		return err
	}

	width, height := GetResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

	bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetDimension(width, height)
	codecContext.SetGopSize(gop)
	codecContext.SetBitRate(bitrate)
	codecContext.SetPixFmt(ist.CodecCtx().PixFmt())

	return nil
}
<commit_msg>removed commented import, simplified getStream<commit_after>package encoders

import (
	"errors"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)

	// create input context
	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	// create output context
	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	//get audio and video stream and the streaMap
	streamMap, srcVideoStream, srcAudioStream, err := getAudioVideoStreamSource(inputCtx, outputCtx, job)
	if err != nil {
		return err
	}
	//calculate total number of frames
	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())
	//process all frames and update the job progress
	err = processAllFramesAndUpdateJobProgress(inputCtx, outputCtx, streamMap, job, dbInstance, totalFrames)
	if err != nil {
		return err
	}

	processNewFrames(inputCtx, outputCtx, streamMap)
	if err != nil {
		return err
	}

	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

func processNewFrames(inputCtx *gmf.FmtCtx, outputCtx *gmf.FmtCtx, streamMap map[int]int) error {
	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		inputStream, err := getStream(inputCtx, 0)
		if err != nil {
			return err
		}
		outputStream, err := getStream(outputCtx, streamMap[inputStream.Index()])
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(outputStream.CodecCtx()); ready {
				configurePacket(p, outputStream, frame)
				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}
			outputStream.Pts++
		}

		gmf.Release(frame)
	}

	return nil
}

func processAllFramesAndUpdateJobProgress(inputCtx *gmf.FmtCtx, outputCtx *gmf.FmtCtx, streamMap map[int]int, job types.Job, dbInstance db.Storage, totalFrames float64) error {
	var lastDelta int64
	for packet := range inputCtx.GetNewPackets() {
		inputStream, err := getStream(inputCtx, packet.StreamIndex())
		if err != nil {
			return err
		}
		outputStream, err := getStream(outputCtx, streamMap[inputStream.Index()])
		if err != nil {
			return err
		}

		framesCount := float64(0)
		for frame := range packet.Frames(inputStream.CodecCtx()) {
			err := proccessFrame(inputStream, outputStream, packet, frame, outputCtx, &lastDelta)
			if err != nil {
				return err
			}

			outputStream.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}

		gmf.Release(packet)
	}
	return nil
}

func getStream(context *gmf.FmtCtx, streamIndex int) (*gmf.Stream, error) {
	return context.GetStream(streamIndex)
}

func getAudioVideoStreamSource(inputCtx *gmf.FmtCtx, outputCtx *gmf.FmtCtx, job types.Job) (map[int]int, *gmf.Stream, *gmf.Stream, error) {
	streamMap := make(map[int]int, 0)

	// add video stream to streamMap
	srcVideoStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)
	if err != nil {
		return nil, nil, nil, errors.New("unable to get the best video stream inside the input context")
	}
	videoCodec := getVideoCodec(job)
	inputIndex, outputIndex, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		return nil, nil, nil, err
	}
	streamMap[inputIndex] = outputIndex

	// add audio stream to streamMap
	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return nil, nil, nil, errors.New("unable to get the best audio stream inside the input context")
	}
	audioCodec := getAudioCodec(job)
	inputIndex, outputIndex, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return nil, nil, nil, err
	}
	streamMap[inputIndex] = outputIndex
	if err := outputCtx.WriteHeader(); err != nil {
		return nil, nil, nil, err
	}

	return streamMap, srcVideoStream, srcAudioStream, nil
}

func configureAudioFrame(packet *gmf.Packet, inputStream *gmf.Stream, outputStream *gmf.Stream, frame *gmf.Frame, lastDelta *int64) {
	fsTb := gmf.AVR{Num: 1, Den: inputStream.CodecCtx().SampleRate()}
	outTb := gmf.AVR{Num: 1, Den: inputStream.CodecCtx().SampleRate()}

	frame.SetPts(packet.Pts())

	pts := gmf.RescaleDelta(inputStream.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), lastDelta, outTb.AVRational())

	frame.SetNbSamples(outputStream.CodecCtx().FrameSize())
	frame.SetFormat(outputStream.CodecCtx().SampleFmt())
	frame.SetChannelLayout(outputStream.CodecCtx().ChannelLayout())
	frame.SetPts(pts)
}

func configurePacket(packet *gmf.Packet, outputStream *gmf.Stream, frame *gmf.Frame) *gmf.Packet {
	if packet.Pts() != gmf.AV_NOPTS_VALUE {
		packet.SetPts(gmf.RescaleQ(packet.Pts(), outputStream.CodecCtx().TimeBase(), outputStream.TimeBase()))
	}

	if packet.Dts() != gmf.AV_NOPTS_VALUE {
		packet.SetDts(gmf.RescaleQ(packet.Dts(), outputStream.CodecCtx().TimeBase(), outputStream.TimeBase()))
	}

	packet.SetStreamIndex(outputStream.Index())

	return packet
}

func proccessFrame(inputStream *gmf.Stream, outputStream *gmf.Stream, packet *gmf.Packet, frame *gmf.Frame, outputCtx *gmf.FmtCtx, lastDelta *int64) error {
	if outputStream.IsAudio() {
		configureAudioFrame(packet, inputStream, outputStream, frame, lastDelta)
	} else {
		frame.SetPts(outputStream.Pts)
	}

	if newPacket, ready, _ := frame.EncodeNewPacket(outputStream.CodecCtx()); ready {
		configurePacket(newPacket, outputStream, frame)
		if err := outputCtx.WritePacket(newPacket); err != nil {
			return err
		}
		gmf.Release(newPacket)
	}

	return nil
}

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, inputStream *gmf.Stream) (int, int, error) {
	var codecContext *gmf.CodecCtx
	var outputStream *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if outputStream = oc.NewStream(codec); outputStream == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(outputStream)

	if codecContext = gmf.NewCodecCtx(codec); codecContext == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(codecContext)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		codecContext.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		codecContext.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		err := setAudioCtxParams(codecContext, inputStream, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		err := setVideoCtxParams(codecContext, inputStream, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if err := codecContext.Open(nil); err != nil {
		return 0, 0, err
	}

	outputStream.SetCodecCtx(codecContext)

	return inputStream.Index(), outputStream.Index(), nil
}

func getProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile != "" {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

func getVideoCodec(job types.Job) string {
	codecs := map[string]string{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
		"aac":    "aac",
	}

	if codec, ok := codecs[job.Preset.Video.Codec]; ok {
		return codec
	}
	return "libx264"
}

func getAudioCodec(job types.Job) string {
	codecs := map[string]string{
		"aac":    "aac",
		"vorbis": "vorbis",
	}
	if codec, ok := codecs[job.Preset.Audio.Codec]; ok {
		return codec
	}
	return "aac"
}

func GetResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}

func setAudioCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetBitRate(bitrate)
	codecContext.SetSampleFmt(ist.CodecCtx().SampleFmt())
	codecContext.SetSampleRate(ist.CodecCtx().SampleRate())
	codecContext.SetChannels(ist.CodecCtx().Channels())
	codecContext.SelectChannelLayout()
	codecContext.SelectSampleRate()
	return nil
}

func setVideoCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	codecContext.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

	if job.Preset.Video.Codec == "h264" {
		profile := getProfile(job)
		codecContext.SetProfile(profile)
	}

	gop, err := strconv.Atoi(job.Preset.Video.GopSize)
	if err != nil {
		return err
	}

	width, height := GetResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

	bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetDimension(width, height)
	codecContext.SetGopSize(gop)
	codecContext.SetBitRate(bitrate)
	codecContext.SetPixFmt(ist.CodecCtx().PixFmt())

	return nil
}
<|endoftext|>


<commit_before>// Copyright ©2016 The gonum 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 native

import (
	"github.com/gonum/blas"
	"github.com/gonum/blas/blas64"
)

// Dsytrd reduces a symmetric n×n matrix A to symmetric tridiagonal form by an
// orthogonal similarity transformation
//  Q^T * A * Q = T
// where Q is an orthonormal matrix and T is symmetric and tridiagonal.
//
// On entry, a contains the elements of the input matrix in the triangle specified
// by uplo. On exit, the diagonal and sub/super-diagonal are overwritten by the
// corresponding elements of the tridiagonal matrix T. The remaining elements in
// the triangle, along with the array tau, contain the data to construct Q as
// the product of elementary reflectors.
//
// If uplo == blas.Upper, Q is constructed with
//  Q = H[n-2] * ... * H[1] * H[0]
// where
//  H[i] = I - tau * v * v^T
// v is constructed as v[i+1:n] = 0, v[i] = 1, v[0:i-1] is stored in A[0:i-1, i+1],
// and tau is in tau[i]. The elements of A are
//  [ d  e  v2 v3 v4]
//  [    d  e  v3 v4]
//  [       d   e v4]
//  [           d  e]
//  [              e]
//
// If uplo == blas.Lower, Q is constructed with
//  Q = H[0] * H[1] * ... * H[n-2]
// where
//  H[i] = I - tau * v * v^T
// v is constructed as v[0:i+1] = 0, v[i+1] = 1, v[i+2:n] is stored in A[i+2:n, i],
// and tau is in tau[i]. The elements of A are
//  [ d            ]
//  [ e  d         ]
//  [v1  e  d      ]
//  [v1 v2  e  d   ]
//  [v1 v2  v3 e  d]
//
// d must have length n, and e and tau must have length n-1. Dsytrd will panic if
// these conditions are not met.
//
// work is temporary storage, and lwork specifies the usable memory length. At minimum,
// lwork >= 1, and Dsytrd will panic otherwise. The amount of blocking is
// limited by the usable length.
// If lwork == -1, instead of computing Dsytrd the optimal work length is stored
// into work[0].
func (impl Implementation) Dsytrd(uplo blas.Uplo, n int, a []float64, lda int, d, e, tau, work []float64, lwork int) {
	upper := uplo == blas.Upper
	opts := "U"
	if !upper {
		opts = "L"
	}
	nb := impl.Ilaenv(1, "DSYTRD", opts, n, -1, -1, -1)
	lworkopt := n * nb
	work[0] = float64(lworkopt)
	if lwork == -1 {
		return
	}
	if n == 0 {
		work[0] = 1
	}
	nx := n

	bi := blas64.Implementation()
	var ldwork int
	if nb > 1 && nb < n {
		// Determine when to cross over from blocked to unblocked code. The last
		// block is always handled by unblocked code.
		opts := "L"
		if upper {
			opts = "U"
		}
		nx = max(nb, impl.Ilaenv(3, "DSYTRD", opts, n, -1, -1, -1))
		if nx < n {
			// Determine if workspace is large enough for blocked code.
			ldwork = nb
			iws := n * ldwork
			if lwork < iws {
				// Not enough workspace to use optimal nb: determine the minimum
				// value of nb and reduce nb or force use of unblocked code by
				// setting nx = n.
				nb = max(lwork/n, 1)
				nbmin := impl.Ilaenv(2, "DSYTRD", opts, n, -1, -1, -1)
				if nb < nbmin {
					nx = n
				}
			}
		} else {
			nx = n
		}
	} else {
		nb = 1
	}
	ldwork = nb

	if upper {
		// Reduce the upper triangle of A. Columns 0:kk are handled by the
		// unblocked method.
		var i int
		kk := n - ((n-nx+nb-1)/nb)*nb
		for i = n - nb; i >= kk; i -= nb {
			// Reduce columns i:i+nb to tridiagonal form and form the matrix W
			// which is needed to update the unreduced part of the matrix.
			impl.Dlatrd(uplo, i+nb, nb, a, lda, e, tau, work, ldwork)

			// Update the unreduced submatrix A[0:i-1,0:i-1], using an update
			// of the form A = A - V*W^T - W*V^T.
			bi.Dsyr2k(uplo, blas.NoTrans, i, nb, -1, a[i:], lda, work, ldwork, 1, a, lda)

			// Copy superdiagonal elements back into A, and diagonal elements into D.
			for j := i; j < i+nb; j++ {
				a[(j-1)*lda+j] = e[j-1]
				d[j] = a[j*lda+j]
			}
		}
		// Use unblocked code to reduce the last or only block
		// check that i == kk.
		impl.Dsytd2(uplo, kk, a, lda, d, e, tau)
	} else {
		var i int
		// Reduce the lower triangle of A.
		for i = 0; i < n-nx; i += nb {
			// Reduce columns 0:i+nb to tridiagonal form and form the matrix W
			// which is needed to update the unreduced part of the matrix.
			impl.Dlatrd(uplo, n-i, nb, a[i*lda+i:], lda, e[i:], tau[i:], work, ldwork)

			// Update the unreduced submatrix A[i+ib:n, i+ib:n], using an update
			// of the form A = A + V*W^T - W*V^T.
			bi.Dsyr2k(uplo, blas.NoTrans, n-i-nb, nb, -1, a[(i+nb)*lda+i:], lda,
				work[nb*ldwork:], ldwork, 1, a[(i+nb)*lda+i+nb:], lda)

			// Copy subdiagonal elements back into A, and diagonal elements into D.
			for j := i; j < i+nb; j++ {
				a[(j+1)*lda+j] = e[j]
				d[j] = a[j*lda+j]
			}
		}
		// Use unblocked code to reduce the last or only block.
		impl.Dsytd2(uplo, n-i, a[i*lda+i:], lda, d[i:], e[i:], tau[i:])
	}
	work[0] = float64(lworkopt)
}
<commit_msg>native: fix missed matrix for typography changes<commit_after>// Copyright ©2016 The gonum 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 native

import (
	"github.com/gonum/blas"
	"github.com/gonum/blas/blas64"
)

// Dsytrd reduces a symmetric n×n matrix A to symmetric tridiagonal form by an
// orthogonal similarity transformation
//  Q^T * A * Q = T
// where Q is an orthonormal matrix and T is symmetric and tridiagonal.
//
// On entry, a contains the elements of the input matrix in the triangle specified
// by uplo. On exit, the diagonal and sub/super-diagonal are overwritten by the
// corresponding elements of the tridiagonal matrix T. The remaining elements in
// the triangle, along with the array tau, contain the data to construct Q as
// the product of elementary reflectors.
//
// If uplo == blas.Upper, Q is constructed with
//  Q = H[n-2] * ... * H[1] * H[0]
// where
//  H[i] = I - tau * v * v^T
// v is constructed as v[i+1:n] = 0, v[i] = 1, v[0:i-1] is stored in A[0:i-1, i+1],
// and tau is in tau[i]. The elements of A are
//  [ d   e  v2  v3  v4]
//  [     d   e  v3  v4]
//  [         d   e  v4]
//  [             d   e]
//  [                 e]
//
// If uplo == blas.Lower, Q is constructed with
//  Q = H[0] * H[1] * ... * H[n-2]
// where
//  H[i] = I - tau * v * v^T
// v is constructed as v[0:i+1] = 0, v[i+1] = 1, v[i+2:n] is stored in A[i+2:n, i],
// and tau is in tau[i]. The elements of A are
//  [ d                ]
//  [ e   d            ]
//  [v1   e   d        ]
//  [v1  v2   e   d    ]
//  [v1  v2  v3   e   d]
//
// d must have length n, and e and tau must have length n-1. Dsytrd will panic if
// these conditions are not met.
//
// work is temporary storage, and lwork specifies the usable memory length. At minimum,
// lwork >= 1, and Dsytrd will panic otherwise. The amount of blocking is
// limited by the usable length.
// If lwork == -1, instead of computing Dsytrd the optimal work length is stored
// into work[0].
func (impl Implementation) Dsytrd(uplo blas.Uplo, n int, a []float64, lda int, d, e, tau, work []float64, lwork int) {
	upper := uplo == blas.Upper
	opts := "U"
	if !upper {
		opts = "L"
	}
	nb := impl.Ilaenv(1, "DSYTRD", opts, n, -1, -1, -1)
	lworkopt := n * nb
	work[0] = float64(lworkopt)
	if lwork == -1 {
		return
	}
	if n == 0 {
		work[0] = 1
	}
	nx := n

	bi := blas64.Implementation()
	var ldwork int
	if nb > 1 && nb < n {
		// Determine when to cross over from blocked to unblocked code. The last
		// block is always handled by unblocked code.
		opts := "L"
		if upper {
			opts = "U"
		}
		nx = max(nb, impl.Ilaenv(3, "DSYTRD", opts, n, -1, -1, -1))
		if nx < n {
			// Determine if workspace is large enough for blocked code.
			ldwork = nb
			iws := n * ldwork
			if lwork < iws {
				// Not enough workspace to use optimal nb: determine the minimum
				// value of nb and reduce nb or force use of unblocked code by
				// setting nx = n.
				nb = max(lwork/n, 1)
				nbmin := impl.Ilaenv(2, "DSYTRD", opts, n, -1, -1, -1)
				if nb < nbmin {
					nx = n
				}
			}
		} else {
			nx = n
		}
	} else {
		nb = 1
	}
	ldwork = nb

	if upper {
		// Reduce the upper triangle of A. Columns 0:kk are handled by the
		// unblocked method.
		var i int
		kk := n - ((n-nx+nb-1)/nb)*nb
		for i = n - nb; i >= kk; i -= nb {
			// Reduce columns i:i+nb to tridiagonal form and form the matrix W
			// which is needed to update the unreduced part of the matrix.
			impl.Dlatrd(uplo, i+nb, nb, a, lda, e, tau, work, ldwork)

			// Update the unreduced submatrix A[0:i-1,0:i-1], using an update
			// of the form A = A - V*W^T - W*V^T.
			bi.Dsyr2k(uplo, blas.NoTrans, i, nb, -1, a[i:], lda, work, ldwork, 1, a, lda)

			// Copy superdiagonal elements back into A, and diagonal elements into D.
			for j := i; j < i+nb; j++ {
				a[(j-1)*lda+j] = e[j-1]
				d[j] = a[j*lda+j]
			}
		}
		// Use unblocked code to reduce the last or only block
		// check that i == kk.
		impl.Dsytd2(uplo, kk, a, lda, d, e, tau)
	} else {
		var i int
		// Reduce the lower triangle of A.
		for i = 0; i < n-nx; i += nb {
			// Reduce columns 0:i+nb to tridiagonal form and form the matrix W
			// which is needed to update the unreduced part of the matrix.
			impl.Dlatrd(uplo, n-i, nb, a[i*lda+i:], lda, e[i:], tau[i:], work, ldwork)

			// Update the unreduced submatrix A[i+ib:n, i+ib:n], using an update
			// of the form A = A + V*W^T - W*V^T.
			bi.Dsyr2k(uplo, blas.NoTrans, n-i-nb, nb, -1, a[(i+nb)*lda+i:], lda,
				work[nb*ldwork:], ldwork, 1, a[(i+nb)*lda+i+nb:], lda)

			// Copy subdiagonal elements back into A, and diagonal elements into D.
			for j := i; j < i+nb; j++ {
				a[(j+1)*lda+j] = e[j]
				d[j] = a[j*lda+j]
			}
		}
		// Use unblocked code to reduce the last or only block.
		impl.Dsytd2(uplo, n-i, a[i*lda+i:], lda, d[i:], e[i:], tau[i:])
	}
	work[0] = float64(lworkopt)
}
<|endoftext|>


<commit_before>package executor

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"testing"
	"time"

	hclog "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/taskenv"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/plugins/drivers"
	tu "github.com/hashicorp/nomad/testutil"
	ps "github.com/mitchellh/go-ps"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var executorFactories = map[string]func(hclog.Logger) Executor{}
var universalFactory = func(l hclog.Logger) Executor {
	return NewExecutor(l)
}

func init() {
	executorFactories["UniversalExecutor"] = universalFactory
}

// testExecutorContext returns an ExecutorContext and AllocDir.
//
// The caller is responsible for calling AllocDir.Destroy() to cleanup.
func testExecutorCommand(t *testing.T) (*ExecCommand, *allocdir.AllocDir) {
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	taskEnv := taskenv.NewBuilder(mock.Node(), alloc, task, "global").Build()

	allocDir := allocdir.NewAllocDir(testlog.HCLogger(t), filepath.Join(os.TempDir(), alloc.ID))
	if err := allocDir.Build(); err != nil {
		t.Fatalf("AllocDir.Build() failed: %v", err)
	}
	if err := allocDir.NewTaskDir(task.Name).Build(false, nil); err != nil {
		allocDir.Destroy()
		t.Fatalf("allocDir.NewTaskDir(%q) failed: %v", task.Name, err)
	}
	td := allocDir.TaskDirs[task.Name]
	cmd := &ExecCommand{
		Env:     taskEnv.List(),
		TaskDir: td.Dir,
		Resources: &drivers.Resources{
			NomadResources: alloc.AllocatedResources.Tasks[task.Name],
		},
	}

	setupRootfs(t, td.Dir)
	configureTLogging(cmd)
	return cmd, allocDir
}

type bufferCloser struct {
	bytes.Buffer
}

func (_ *bufferCloser) Close() error { return nil }

func configureTLogging(cmd *ExecCommand) (stdout bufferCloser, stderr bufferCloser) {
	cmd.SetWriters(&stdout, &stderr)
	return
}

func TestExecutor_Start_Invalid(pt *testing.T) {
	pt.Parallel()
	invalid := "/bin/foobar"
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = invalid
			execCmd.Args = []string{"1"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			_, err := executor.Launch(execCmd)
			require.Error(err)
		})
	}
}

func TestExecutor_Start_Wait_Failure_Code(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/date"
			execCmd.Args = []string{"fail"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)
			ps, _ = executor.Wait(context.Background())
			require.NotZero(ps.ExitCode, "expected exit code to be non zero")
			require.NoError(executor.Shutdown("SIGINT", 100*time.Millisecond))
		})
	}
}

func TestExecutor_Start_Wait(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/echo"
			execCmd.Args = []string{"hello world"}

			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)

			ps, err = executor.Wait(context.Background())
			require.NoError(err)
			require.NoError(executor.Shutdown("SIGINT", 100*time.Millisecond))

			expected := "hello world"
			tu.WaitForResult(func() (bool, error) {
				outWriter, _ := execCmd.GetWriters()
				output := outWriter.(*bufferCloser).String()
				act := strings.TrimSpace(string(output))
				if expected != act {
					return false, fmt.Errorf("expected: '%s' actual: '%s'", expected, act)
				}
				return true, nil
			}, func(err error) {
				require.NoError(err)
			})
		})
	}
}

func TestExecutor_WaitExitSignal(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/sleep"
			execCmd.Args = []string{"10000"}
			execCmd.ResourceLimits = true
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)

			go func() {
				tu.WaitForResult(func() (bool, error) {
					ch, err := executor.Stats(context.Background(), time.Second)
					if err != nil {
						return false, err
					}
					select {
					case <-time.After(time.Second):
						return false, fmt.Errorf("stats failed to send on interval")
					case ru := <-ch:
						assert.NotEmpty(t, ru.Pids, "no pids recorded in stats")
						assert.NotZero(t, ru.ResourceUsage.MemoryStats.RSS)
						assert.WithinDuration(t, time.Now(), time.Unix(0, ru.Timestamp), time.Second)
					}
					proc, err := os.FindProcess(ps.Pid)
					if err != nil {
						return false, err
					}
					err = proc.Signal(syscall.SIGKILL)
					if err != nil {
						return false, err
					}
					return true, nil
				}, func(err error) {
					assert.NoError(t, executor.Signal(os.Kill))
					assert.NoError(t, err)
				})
			}()

			ps, err = executor.Wait(context.Background())
			require.NoError(err)
			require.Equal(ps.Signal, int(syscall.SIGKILL))
		})
	}
}

func TestExecutor_Start_Kill(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/sleep"
			execCmd.Args = []string{"10"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)

			require.NoError(executor.Shutdown("SIGINT", 100*time.Millisecond))

			time.Sleep(time.Duration(tu.TestMultiplier()*2) * time.Second)
			outWriter, _ := execCmd.GetWriters()
			output := outWriter.(*bufferCloser).String()
			expected := ""
			act := strings.TrimSpace(string(output))
			if act != expected {
				t.Fatalf("Command output incorrectly: want %v; got %v", expected, act)
			}
		})
	}
}

func TestExecutor_Shutdown_Exit(t *testing.T) {
	require := require.New(t)
	t.Parallel()
	execCmd, allocDir := testExecutorCommand(t)
	execCmd.Cmd = "/bin/sleep"
	execCmd.Args = []string{"100"}
	cfg := &ExecutorConfig{
		LogFile: "/dev/null",
	}
	executor, pluginClient, err := CreateExecutor(testlog.HCLogger(t), nil, cfg)
	require.NoError(err)

	proc, err := executor.Launch(execCmd)
	require.NoError(err)
	require.NotZero(proc.Pid)

	executor.Shutdown("", 0)
	pluginClient.Kill()
	tu.WaitForResult(func() (bool, error) {
		p, err := ps.FindProcess(proc.Pid)
		if err != nil {
			return false, err
		}
		return p == nil, fmt.Errorf("process found: %d", proc.Pid)
	}, func(err error) {
		require.NoError(err)
	})
	require.NoError(allocDir.Destroy())
}

func TestUniversalExecutor_MakeExecutable(t *testing.T) {
	t.Parallel()
	// Create a temp file
	f, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	defer os.Remove(f.Name())

	// Set its permissions to be non-executable
	f.Chmod(os.FileMode(0610))

	err = makeExecutable(f.Name())
	if err != nil {
		t.Fatalf("makeExecutable() failed: %v", err)
	}

	// Check the permissions
	stat, err := f.Stat()
	if err != nil {
		t.Fatalf("Stat() failed: %v", err)
	}

	act := stat.Mode().Perm()
	exp := os.FileMode(0755)
	if act != exp {
		t.Fatalf("expected permissions %v; got %v", exp, act)
	}
}

func TestUniversalExecutor_LookupPath(t *testing.T) {
	t.Parallel()
	// Create a temp dir
	tmpDir, err := ioutil.TempDir("", "")
	require := require.New(t)
	require.Nil(err)
	defer os.Remove(tmpDir)

	// Make a foo subdir
	os.MkdirAll(filepath.Join(tmpDir, "foo"), 0700)

	// Write a file under foo
	filePath := filepath.Join(tmpDir, "foo", "tmp.txt")
	err = ioutil.WriteFile(filePath, []byte{1, 2}, os.ModeAppend)
	require.Nil(err)

	// Lookup with full path to binary
	_, err = lookupBin("dummy", filePath)
	require.Nil(err)

	// Write a file under local subdir
	os.MkdirAll(filepath.Join(tmpDir, "local"), 0700)
	filePath2 := filepath.Join(tmpDir, "local", "tmp.txt")
	ioutil.WriteFile(filePath2, []byte{1, 2}, os.ModeAppend)

	// Lookup with file name, should find the one we wrote above
	_, err = lookupBin(tmpDir, "tmp.txt")
	require.Nil(err)

	// Write a file under task dir
	filePath3 := filepath.Join(tmpDir, "tmp.txt")
	ioutil.WriteFile(filePath3, []byte{1, 2}, os.ModeAppend)

	// Lookup with file name, should find the one we wrote above
	_, err = lookupBin(tmpDir, "tmp.txt")
	require.Nil(err)

}

// setupRoootfs setups the rootfs for libcontainer executor
// It uses busybox to make some binaries available - somewhat cheaper
// than mounting the underlying host filesystem
func setupRootfs(t *testing.T, rootfs string) {
	paths := []string{
		"/bin/sh",
		"/bin/sleep",
		"/bin/echo",
		"/bin/date",
	}

	for _, p := range paths {
		setupRootfsBinary(t, rootfs, p)
	}
}

// setupRootfsBinary installs a busybox link in the desired path
func setupRootfsBinary(t *testing.T, rootfs, path string) {
	t.Helper()

	dst := filepath.Join(rootfs, path)
	err := os.MkdirAll(filepath.Dir(dst), 666)
	require.NoError(t, err)

	src := filepath.Join(
		"test-resources", "busybox",
		fmt.Sprintf("busybox-%s", runtime.GOARCH),
	)

	err = os.Link(src, dst)
	require.NoError(t, err)
}
<commit_msg>test kill wait<commit_after>package executor

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"testing"
	"time"

	hclog "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/taskenv"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/plugins/drivers"
	tu "github.com/hashicorp/nomad/testutil"
	ps "github.com/mitchellh/go-ps"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var executorFactories = map[string]func(hclog.Logger) Executor{}
var universalFactory = func(l hclog.Logger) Executor {
	return NewExecutor(l)
}

func init() {
	executorFactories["UniversalExecutor"] = universalFactory
}

// testExecutorContext returns an ExecutorContext and AllocDir.
//
// The caller is responsible for calling AllocDir.Destroy() to cleanup.
func testExecutorCommand(t *testing.T) (*ExecCommand, *allocdir.AllocDir) {
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	taskEnv := taskenv.NewBuilder(mock.Node(), alloc, task, "global").Build()

	allocDir := allocdir.NewAllocDir(testlog.HCLogger(t), filepath.Join(os.TempDir(), alloc.ID))
	if err := allocDir.Build(); err != nil {
		t.Fatalf("AllocDir.Build() failed: %v", err)
	}
	if err := allocDir.NewTaskDir(task.Name).Build(false, nil); err != nil {
		allocDir.Destroy()
		t.Fatalf("allocDir.NewTaskDir(%q) failed: %v", task.Name, err)
	}
	td := allocDir.TaskDirs[task.Name]
	cmd := &ExecCommand{
		Env:     taskEnv.List(),
		TaskDir: td.Dir,
		Resources: &drivers.Resources{
			NomadResources: alloc.AllocatedResources.Tasks[task.Name],
		},
	}

	setupRootfs(t, td.Dir)
	configureTLogging(cmd)
	return cmd, allocDir
}

type bufferCloser struct {
	bytes.Buffer
}

func (_ *bufferCloser) Close() error { return nil }

func configureTLogging(cmd *ExecCommand) (stdout bufferCloser, stderr bufferCloser) {
	cmd.SetWriters(&stdout, &stderr)
	return
}

func TestExecutor_Start_Invalid(pt *testing.T) {
	pt.Parallel()
	invalid := "/bin/foobar"
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = invalid
			execCmd.Args = []string{"1"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			_, err := executor.Launch(execCmd)
			require.Error(err)
		})
	}
}

func TestExecutor_Start_Wait_Failure_Code(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/date"
			execCmd.Args = []string{"fail"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)
			ps, _ = executor.Wait(context.Background())
			require.NotZero(ps.ExitCode, "expected exit code to be non zero")
			require.NoError(executor.Shutdown("SIGINT", 100*time.Millisecond))
		})
	}
}

func TestExecutor_Start_Wait(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/echo"
			execCmd.Args = []string{"hello world"}

			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)

			ps, err = executor.Wait(context.Background())
			require.NoError(err)
			require.NoError(executor.Shutdown("SIGINT", 100*time.Millisecond))

			expected := "hello world"
			tu.WaitForResult(func() (bool, error) {
				outWriter, _ := execCmd.GetWriters()
				output := outWriter.(*bufferCloser).String()
				act := strings.TrimSpace(string(output))
				if expected != act {
					return false, fmt.Errorf("expected: '%s' actual: '%s'", expected, act)
				}
				return true, nil
			}, func(err error) {
				require.NoError(err)
			})
		})
	}
}

func TestExecutor_WaitExitSignal(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/sleep"
			execCmd.Args = []string{"10000"}
			execCmd.ResourceLimits = true
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)

			go func() {
				tu.WaitForResult(func() (bool, error) {
					ch, err := executor.Stats(context.Background(), time.Second)
					if err != nil {
						return false, err
					}
					select {
					case <-time.After(time.Second):
						return false, fmt.Errorf("stats failed to send on interval")
					case ru := <-ch:
						assert.NotEmpty(t, ru.Pids, "no pids recorded in stats")
						assert.NotZero(t, ru.ResourceUsage.MemoryStats.RSS)
						assert.WithinDuration(t, time.Now(), time.Unix(0, ru.Timestamp), time.Second)
					}
					proc, err := os.FindProcess(ps.Pid)
					if err != nil {
						return false, err
					}
					err = proc.Signal(syscall.SIGKILL)
					if err != nil {
						return false, err
					}
					return true, nil
				}, func(err error) {
					assert.NoError(t, executor.Signal(os.Kill))
					assert.NoError(t, err)
				})
			}()

			ps, err = executor.Wait(context.Background())
			require.NoError(err)
			require.Equal(ps.Signal, int(syscall.SIGKILL))
		})
	}
}

func TestExecutor_Start_Kill(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/sleep"
			execCmd.Args = []string{"10"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)

			require.NoError(executor.Shutdown("SIGINT", 100*time.Millisecond))

			time.Sleep(time.Duration(tu.TestMultiplier()*2) * time.Second)
			outWriter, _ := execCmd.GetWriters()
			output := outWriter.(*bufferCloser).String()
			expected := ""
			act := strings.TrimSpace(string(output))
			if act != expected {
				t.Fatalf("Command output incorrectly: want %v; got %v", expected, act)
			}
		})
	}
}

func TestExecutor_Shutdown_Exit(t *testing.T) {
	require := require.New(t)
	t.Parallel()
	execCmd, allocDir := testExecutorCommand(t)
	execCmd.Cmd = "/bin/sleep"
	execCmd.Args = []string{"100"}
	cfg := &ExecutorConfig{
		LogFile: "/dev/null",
	}
	executor, pluginClient, err := CreateExecutor(testlog.HCLogger(t), nil, cfg)
	require.NoError(err)

	proc, err := executor.Launch(execCmd)
	require.NoError(err)
	require.NotZero(proc.Pid)

	executor.Shutdown("", 0)
	pluginClient.Kill()
	tu.WaitForResult(func() (bool, error) {
		p, err := ps.FindProcess(proc.Pid)
		if err != nil {
			return false, err
		}
		return p == nil, fmt.Errorf("process found: %d", proc.Pid)
	}, func(err error) {
		require.NoError(err)
	})
	require.NoError(allocDir.Destroy())
}

func TestUniversalExecutor_MakeExecutable(t *testing.T) {
	t.Parallel()
	// Create a temp file
	f, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	defer os.Remove(f.Name())

	// Set its permissions to be non-executable
	f.Chmod(os.FileMode(0610))

	err = makeExecutable(f.Name())
	if err != nil {
		t.Fatalf("makeExecutable() failed: %v", err)
	}

	// Check the permissions
	stat, err := f.Stat()
	if err != nil {
		t.Fatalf("Stat() failed: %v", err)
	}

	act := stat.Mode().Perm()
	exp := os.FileMode(0755)
	if act != exp {
		t.Fatalf("expected permissions %v; got %v", exp, act)
	}
}

func TestUniversalExecutor_LookupPath(t *testing.T) {
	t.Parallel()
	// Create a temp dir
	tmpDir, err := ioutil.TempDir("", "")
	require := require.New(t)
	require.Nil(err)
	defer os.Remove(tmpDir)

	// Make a foo subdir
	os.MkdirAll(filepath.Join(tmpDir, "foo"), 0700)

	// Write a file under foo
	filePath := filepath.Join(tmpDir, "foo", "tmp.txt")
	err = ioutil.WriteFile(filePath, []byte{1, 2}, os.ModeAppend)
	require.Nil(err)

	// Lookup with full path to binary
	_, err = lookupBin("dummy", filePath)
	require.Nil(err)

	// Write a file under local subdir
	os.MkdirAll(filepath.Join(tmpDir, "local"), 0700)
	filePath2 := filepath.Join(tmpDir, "local", "tmp.txt")
	ioutil.WriteFile(filePath2, []byte{1, 2}, os.ModeAppend)

	// Lookup with file name, should find the one we wrote above
	_, err = lookupBin(tmpDir, "tmp.txt")
	require.Nil(err)

	// Write a file under task dir
	filePath3 := filepath.Join(tmpDir, "tmp.txt")
	ioutil.WriteFile(filePath3, []byte{1, 2}, os.ModeAppend)

	// Lookup with file name, should find the one we wrote above
	_, err = lookupBin(tmpDir, "tmp.txt")
	require.Nil(err)

}

// setupRoootfs setups the rootfs for libcontainer executor
// It uses busybox to make some binaries available - somewhat cheaper
// than mounting the underlying host filesystem
func setupRootfs(t *testing.T, rootfs string) {
	paths := []string{
		"/bin/sh",
		"/bin/sleep",
		"/bin/echo",
		"/bin/date",
	}

	for _, p := range paths {
		setupRootfsBinary(t, rootfs, p)
	}
}

// setupRootfsBinary installs a busybox link in the desired path
func setupRootfsBinary(t *testing.T, rootfs, path string) {
	t.Helper()

	dst := filepath.Join(rootfs, path)
	err := os.MkdirAll(filepath.Dir(dst), 666)
	require.NoError(t, err)

	src := filepath.Join(
		"test-resources", "busybox",
		fmt.Sprintf("busybox-%s", runtime.GOARCH),
	)

	err = os.Link(src, dst)
	require.NoError(t, err)
}

func TestExecutor_Start_Kill_Immediately_NoGrace(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/sleep"
			execCmd.Args = []string{"100"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)

			waitCh := make(chan interface{})
			go func() {
				defer close(waitCh)
				executor.Wait(context.Background())
			}()

			require.NoError(executor.Shutdown("SIGKILL", 0))

			select {
			case <-waitCh:
				// all good!
			case <-time.After(4 * time.Second * time.Duration(tu.TestMultiplier())):
				require.Fail("process did not terminate despite SIGKILL")
			}
		})
	}
}

func TestExecutor_Start_Kill_Immediately_WithGrace(pt *testing.T) {
	pt.Parallel()
	for name, factory := range executorFactories {
		pt.Run(name, func(t *testing.T) {
			require := require.New(t)
			execCmd, allocDir := testExecutorCommand(t)
			execCmd.Cmd = "/bin/sleep"
			execCmd.Args = []string{"100"}
			defer allocDir.Destroy()
			executor := factory(testlog.HCLogger(t))
			defer executor.Shutdown("", 0)

			ps, err := executor.Launch(execCmd)
			require.NoError(err)
			require.NotZero(ps.Pid)

			waitCh := make(chan interface{})
			go func() {
				defer close(waitCh)
				executor.Wait(context.Background())
			}()

			require.NoError(executor.Shutdown("SIGKILL", 100*time.Millisecond))

			select {
			case <-waitCh:
				// all good!
			case <-time.After(4 * time.Second * time.Duration(tu.TestMultiplier())):
				require.Fail("process did not terminate despite SIGKILL")
			}
		})
	}
}
<|endoftext|>


<commit_before>/*
This file implements the uWSGI protocol.
This implements run as net.Listener:

		l, err = net.Listen("unix", "/path/to/socket")
		http.Serve(&UwsgiListener{l}, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			w.Header().Set("Content-Length", 11)
			w.Write([]byte("hello world"))
		})
*/

package uwsgi

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Listener struct {
	net.Listener
}

type Conn struct {
	net.Conn
	env     map[string][]string
	reader  io.Reader
	hdrdone bool
	ready   bool
	readych chan bool
	err     error
}

func (c *Conn) Read(b []byte) (n int, e error) {
	// Wait until headers have been processed
	if !c.ready && c.err == nil {
		<-c.readych
		c.ready = true
	}
	if c.err != nil {
		return 0, c.err
	}

	// After headers have been read by HTTP server, transfer
	// socket over to the underlying connection for direct read.
	if !c.hdrdone {
		n, e = c.reader.Read(b)
		if n == 0 || e != nil {
			c.hdrdone = true
		}
	}
	if c.hdrdone {
		n, e = c.Conn.Read(b)
		c.err = e
	}

	return n, e
}

func (c *Conn) Write(b []byte) (int, error) {
	if c.err != nil {
		return 0, c.err
	}

	return c.Conn.Write(b)
}

func (c *Conn) SetDeadline(t time.Time) error {
	if c.err != nil {
		return c.err
	}

	return c.Conn.SetDeadline(t)
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	if c.err != nil {
		return c.err
	}

	return c.Conn.SetReadDeadline(t)
}

func (c *Conn) SetWriteDeadline(t time.Time) error {
	if c.err != nil {
		return c.err
	}

	return c.Conn.SetWriteDeadline(t)
}

var headerMappings = map[string]string{
	"HTTP_HOST":              "Host",
	"CONTENT_TYPE":           "Content-Type",
	"HTTP_ACCEPT":            "Accept",
	"HTTP_ACCEPT_ENCODING":   "Accept-Encoding",
	"HTTP_ACCEPT_LANGUAGE":   "Accept-Language",
	"HTTP_ACCEPT_CHARSET":    "Accept-Charset",
	"HTTP_CONTENT_TYPE":      "Content-Type",
	"HTTP_COOKIE":            "Cookie",
	"HTTP_IF_MATCH":          "If-Match",
	"HTTP_IF_MODIFIED_SINCE": "If-Modified-Since",
	"HTTP_IF_NONE_MATCH":     "If-None-Match",
	"HTTP_IF_RANGE":          "If-Range",
	"HTTP_RANGE":             "Range",
	"HTTP_REFERER":           "Referer",
	"HTTP_USER_AGENT":        "User-Agent",
	"HTTP_X_REQUESTED_WITH":  "Requested-With",
}

// Accept conduct as net.Listener. uWSGI protocol is working good for CGI.
// This function parse headers and pass to the Server.
func (l *Listener) Accept() (net.Conn, error) {
	fd, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	c := &Conn{fd, make(map[string][]string), buf, false, false, make(chan bool, 1), nil}

	go func() {
		/*
		 * uwsgi header:
		 * struct {
		 *    uint8  modifier1;
		 *    uint16 datasize;
		 *    uint8  modifier2;
		 * }
		 *  -- for HTTP, mod1 and mod2 = 0
		 */
		var head [4]byte
		fd.Read(head[:])
		b := []byte{head[1], head[2]}
		envsize := binary.LittleEndian.Uint16(b)

		envbuf := make([]byte, envsize)
		if _, err := io.ReadFull(fd, envbuf); err != nil {
			fd.Close()
			c.err = err
			return
		}

		/*
		 * uwsgi vars are linear lists of the form:
		 * struct {
		 *   uint16 key_size;
		 *   uint8  key[key_size];
		 *   uint16 val_size;
		 *   uint8  val[val_size];
		 * }
		 */
		i := uint16(0)
		var reqMethod string
		var reqURI string
		var reqProtocol string
		for {
			// Ensure no corrupted payload; shouldn't happen but it has...
			if i+1 >= uint16(len(envbuf)) {
				break
			}
			b := []byte{envbuf[i], envbuf[i+1]}
			kl := binary.LittleEndian.Uint16(b)
			i += 2

			if i+kl > uint16(len(envbuf)) {
				fd.Close()
				c.err = errors.New("Invalid uwsgi request; uwsgi vars index out of range")
				return
			}

			k := string(envbuf[i : i+kl])
			i += kl

			if i+1 >= uint16(len(envbuf)) {
				fd.Close()
				c.err = errors.New("Invalid uwsgi request; uwsgi vars index out of range")
				return
			}

			b = []byte{envbuf[i], envbuf[i+1]}
			vl := binary.LittleEndian.Uint16(b)
			i += 2

			if i+vl > uint16(len(envbuf)) {
				fd.Close()
				c.err = errors.New("Invalid uwsgi request; uwsgi vars index out of range")
				return
			}

			v := string(envbuf[i : i+vl])
			i += vl

			if k == "REQUEST_METHOD" {
				reqMethod = v
			} else if k == "REQUEST_URI" {
				reqURI = v
			} else if k == "SERVER_PROTOCOL" {
				reqProtocol = v
			}

			val, ok := c.env[k]
			if !ok {
				val = make([]string, 0, 2)
			}
			val = append(val, v)
			c.env[k] = val

			if i >= envsize {
				break
			}
		}

		if reqProtocol == "" {
			// Invalid protocol
			fd.Close()
			c.err = errors.New("Invalid uwsgi request; no protocol specified")
			return
		}

		fmt.Fprintf(buf, "%s %s %s\r\n", reqMethod, reqURI, reqProtocol)

		var cl int64
		for i := range c.env {
			switch i {
			case "CONTENT_LENGTH":
				cl, _ = strconv.ParseInt(c.env[i][0], 10, 64)
				if cl > 0 {
					fmt.Fprintf(buf, "Content-Length: %d\r\n", cl)
				}
			default:
				hname, ok := headerMappings[i]
				if !ok {
					hname = i
				}
				for v := range c.env[i] {
					fmt.Fprintf(buf, "%s: %s\r\n", hname, c.env[i][v])
				}
			}
		}

		buf.Write([]byte("\r\n"))

		// Signal to indicate header processing is complete and remaining
		// payload can be read from the socket itself.
		c.readych <- true
	}()

	return c, nil
}

type Passenger struct {
	Net  string
	Addr string
}

var trailingPort = regexp.MustCompile(`:([0-9]+)$`)

func (p Passenger) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	conn, err := net.Dial(p.Net, p.Addr)
	if err != nil {
		panic(err.Error())
	}
	defer conn.Close()

	port := "80"
	if matches := trailingPort.FindStringSubmatch(req.Host); len(matches) != 0 {
		port = matches[1]
	}

	header := make(map[string][]string)
	header["REQUEST_METHOD"] = []string{req.Method}
	header["REQUEST_URI"] = []string{req.RequestURI}
	header["CONTENT_LENGTH"] = []string{strconv.Itoa(int(req.ContentLength))}
	header["SERVER_PROTOCOL"] = []string{req.Proto}
	header["SERVER_NAME"] = []string{req.Host}
	header["SERVER_ADDR"] = []string{req.RemoteAddr}
	header["SERVER_PORT"] = []string{port}
	header["REMOTE_HOST"] = []string{req.RemoteAddr}
	header["REMOTE_ADDR"] = []string{req.RemoteAddr}
	header["SCRIPT_NAME"] = []string{req.URL.Path}
	header["PATH_INFO"] = []string{req.URL.Path}
	header["QUERY_STRING"] = []string{req.URL.RawQuery}
	if ctype := req.Header.Get("Content-Type"); ctype != "" {
		header["CONTENT_TYPE"] = []string{ctype}
	}
	for k, v := range req.Header {
		if _, ok := header[k]; ok == false {
			k = "HTTP_" + strings.ToUpper(strings.Replace(k, "-", "_", -1))
			header[k] = v
		}
	}

	var size uint16
	for k, v := range header {
		for _, vv := range v {
			size += uint16(len(([]byte)(k))) + 2
			size += uint16(len(([]byte)(vv))) + 2
		}
	}

	hsize := make([]byte, 4)
	binary.LittleEndian.PutUint16(hsize[1:3], size)
	conn.Write(hsize)

	for k, v := range header {
		for _, vv := range v {
			binary.Write(conn, binary.LittleEndian, uint16(len(([]byte)(k))))
			conn.Write([]byte(k))
			binary.Write(conn, binary.LittleEndian, uint16(len(([]byte)(vv))))
			conn.Write([]byte(vv))
		}
	}

	io.Copy(conn, req.Body)

	res, err := http.ReadResponse(bufio.NewReader(conn), req)
	if err != nil {
		panic(err.Error())
	}
	for k, v := range res.Header {
		w.Header().Del(k)
		for _, vv := range v {
			w.Header().Add(k, vv)
		}
	}
	io.Copy(w, res.Body)
}
<commit_msg>golint<commit_after>/*
This file implements the uWSGI protocol.
This implements run as net.Listener:

		l, err = net.Listen("unix", "/path/to/socket")
		http.Serve(&UwsgiListener{l}, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			w.Header().Set("Content-Length", 11)
			w.Write([]byte("hello world"))
		})
*/

package uwsgi

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// Listener behave as net.Listener
type Listener struct {
	net.Listener
}

// Conn is connection for uWSGI
type Conn struct {
	net.Conn
	env     map[string][]string
	reader  io.Reader
	hdrdone bool
	ready   bool
	readych chan bool
	err     error
}

func (c *Conn) Read(b []byte) (n int, e error) {
	// Wait until headers have been processed
	if !c.ready && c.err == nil {
		<-c.readych
		c.ready = true
	}
	if c.err != nil {
		return 0, c.err
	}

	// After headers have been read by HTTP server, transfer
	// socket over to the underlying connection for direct read.
	if !c.hdrdone {
		n, e = c.reader.Read(b)
		if n == 0 || e != nil {
			c.hdrdone = true
		}
	}
	if c.hdrdone {
		n, e = c.Conn.Read(b)
		c.err = e
	}

	return n, e
}

// Writer behave as same as net.Listener
func (c *Conn) Write(b []byte) (int, error) {
	if c.err != nil {
		return 0, c.err
	}

	return c.Conn.Write(b)
}

// SetDeadline behave as same as net.Listener
func (c *Conn) SetDeadline(t time.Time) error {
	if c.err != nil {
		return c.err
	}

	return c.Conn.SetDeadline(t)
}

// SetReadDeadline behave as same as net.Listener
func (c *Conn) SetReadDeadline(t time.Time) error {
	if c.err != nil {
		return c.err
	}

	return c.Conn.SetReadDeadline(t)
}

// SetWriteDeadline behave as same as net.Listener
func (c *Conn) SetWriteDeadline(t time.Time) error {
	if c.err != nil {
		return c.err
	}

	return c.Conn.SetWriteDeadline(t)
}

var headerMappings = map[string]string{
	"HTTP_HOST":              "Host",
	"CONTENT_TYPE":           "Content-Type",
	"HTTP_ACCEPT":            "Accept",
	"HTTP_ACCEPT_ENCODING":   "Accept-Encoding",
	"HTTP_ACCEPT_LANGUAGE":   "Accept-Language",
	"HTTP_ACCEPT_CHARSET":    "Accept-Charset",
	"HTTP_CONTENT_TYPE":      "Content-Type",
	"HTTP_COOKIE":            "Cookie",
	"HTTP_IF_MATCH":          "If-Match",
	"HTTP_IF_MODIFIED_SINCE": "If-Modified-Since",
	"HTTP_IF_NONE_MATCH":     "If-None-Match",
	"HTTP_IF_RANGE":          "If-Range",
	"HTTP_RANGE":             "Range",
	"HTTP_REFERER":           "Referer",
	"HTTP_USER_AGENT":        "User-Agent",
	"HTTP_X_REQUESTED_WITH":  "Requested-With",
}

// Accept conduct as net.Listener. uWSGI protocol is working good for CGI.
// This function parse headers and pass to the Server.
func (l *Listener) Accept() (net.Conn, error) {
	fd, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	c := &Conn{fd, make(map[string][]string), buf, false, false, make(chan bool, 1), nil}

	go func() {
		/*
		 * uwsgi header:
		 * struct {
		 *    uint8  modifier1;
		 *    uint16 datasize;
		 *    uint8  modifier2;
		 * }
		 *  -- for HTTP, mod1 and mod2 = 0
		 */
		var head [4]byte
		fd.Read(head[:])
		b := []byte{head[1], head[2]}
		envsize := binary.LittleEndian.Uint16(b)

		envbuf := make([]byte, envsize)
		if _, err := io.ReadFull(fd, envbuf); err != nil {
			fd.Close()
			c.err = err
			return
		}

		/*
		 * uwsgi vars are linear lists of the form:
		 * struct {
		 *   uint16 key_size;
		 *   uint8  key[key_size];
		 *   uint16 val_size;
		 *   uint8  val[val_size];
		 * }
		 */
		i := uint16(0)
		var reqMethod string
		var reqURI string
		var reqProtocol string
		for {
			// Ensure no corrupted payload; shouldn't happen but it has...
			if i+1 >= uint16(len(envbuf)) {
				break
			}
			b := []byte{envbuf[i], envbuf[i+1]}
			kl := binary.LittleEndian.Uint16(b)
			i += 2

			if i+kl > uint16(len(envbuf)) {
				fd.Close()
				c.err = errors.New("Invalid uwsgi request; uwsgi vars index out of range")
				return
			}

			k := string(envbuf[i : i+kl])
			i += kl

			if i+1 >= uint16(len(envbuf)) {
				fd.Close()
				c.err = errors.New("Invalid uwsgi request; uwsgi vars index out of range")
				return
			}

			b = []byte{envbuf[i], envbuf[i+1]}
			vl := binary.LittleEndian.Uint16(b)
			i += 2

			if i+vl > uint16(len(envbuf)) {
				fd.Close()
				c.err = errors.New("Invalid uwsgi request; uwsgi vars index out of range")
				return
			}

			v := string(envbuf[i : i+vl])
			i += vl

			if k == "REQUEST_METHOD" {
				reqMethod = v
			} else if k == "REQUEST_URI" {
				reqURI = v
			} else if k == "SERVER_PROTOCOL" {
				reqProtocol = v
			}

			val, ok := c.env[k]
			if !ok {
				val = make([]string, 0, 2)
			}
			val = append(val, v)
			c.env[k] = val

			if i >= envsize {
				break
			}
		}

		if reqProtocol == "" {
			// Invalid protocol
			fd.Close()
			c.err = errors.New("Invalid uwsgi request; no protocol specified")
			return
		}

		fmt.Fprintf(buf, "%s %s %s\r\n", reqMethod, reqURI, reqProtocol)

		var cl int64
		for i := range c.env {
			switch i {
			case "CONTENT_LENGTH":
				cl, _ = strconv.ParseInt(c.env[i][0], 10, 64)
				if cl > 0 {
					fmt.Fprintf(buf, "Content-Length: %d\r\n", cl)
				}
			default:
				hname, ok := headerMappings[i]
				if !ok {
					hname = i
				}
				for v := range c.env[i] {
					fmt.Fprintf(buf, "%s: %s\r\n", hname, c.env[i][v])
				}
			}
		}

		buf.Write([]byte("\r\n"))

		// Signal to indicate header processing is complete and remaining
		// payload can be read from the socket itself.
		c.readych <- true
	}()

	return c, nil
}

// Passenger works as uWSGI transport
type Passenger struct {
	Net  string
	Addr string
}

var trailingPort = regexp.MustCompile(`:([0-9]+)$`)

func (p Passenger) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	conn, err := net.Dial(p.Net, p.Addr)
	if err != nil {
		panic(err.Error())
	}
	defer conn.Close()

	port := "80"
	if matches := trailingPort.FindStringSubmatch(req.Host); len(matches) != 0 {
		port = matches[1]
	}

	header := make(map[string][]string)
	header["REQUEST_METHOD"] = []string{req.Method}
	header["REQUEST_URI"] = []string{req.RequestURI}
	header["CONTENT_LENGTH"] = []string{strconv.Itoa(int(req.ContentLength))}
	header["SERVER_PROTOCOL"] = []string{req.Proto}
	header["SERVER_NAME"] = []string{req.Host}
	header["SERVER_ADDR"] = []string{req.RemoteAddr}
	header["SERVER_PORT"] = []string{port}
	header["REMOTE_HOST"] = []string{req.RemoteAddr}
	header["REMOTE_ADDR"] = []string{req.RemoteAddr}
	header["SCRIPT_NAME"] = []string{req.URL.Path}
	header["PATH_INFO"] = []string{req.URL.Path}
	header["QUERY_STRING"] = []string{req.URL.RawQuery}
	if ctype := req.Header.Get("Content-Type"); ctype != "" {
		header["CONTENT_TYPE"] = []string{ctype}
	}
	for k, v := range req.Header {
		if _, ok := header[k]; ok == false {
			k = "HTTP_" + strings.ToUpper(strings.Replace(k, "-", "_", -1))
			header[k] = v
		}
	}

	var size uint16
	for k, v := range header {
		for _, vv := range v {
			size += uint16(len(([]byte)(k))) + 2
			size += uint16(len(([]byte)(vv))) + 2
		}
	}

	hsize := make([]byte, 4)
	binary.LittleEndian.PutUint16(hsize[1:3], size)
	conn.Write(hsize)

	for k, v := range header {
		for _, vv := range v {
			binary.Write(conn, binary.LittleEndian, uint16(len(([]byte)(k))))
			conn.Write([]byte(k))
			binary.Write(conn, binary.LittleEndian, uint16(len(([]byte)(vv))))
			conn.Write([]byte(vv))
		}
	}

	io.Copy(conn, req.Body)

	res, err := http.ReadResponse(bufio.NewReader(conn), req)
	if err != nil {
		panic(err.Error())
	}
	for k, v := range res.Header {
		w.Header().Del(k)
		for _, vv := range v {
			w.Header().Add(k, vv)
		}
	}
	io.Copy(w, res.Body)
}
<|endoftext|>


<commit_before>/*
The MIT License (MIT)

Copyright (c) 2016 winlin

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

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

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

/*
 This the main entrance of https-proxy, proxy to api or other http server.
*/
package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/ossrs/go-oryx-lib/https"
	"net/http"
	"os"
	"strings"
	"sync"
)

func main() {
	var httpPort, httpsPort int
	flag.IntVar(&httpPort, "http", 80, "http listen at. 0 to disable http.")
	flag.IntVar(&httpsPort, "https", 443, "https listen at. 0 to disable https. 443 to serve.")
	flag.Parse()

	if httpsPort != 0 && httpsPort != 443 {
		fmt.Println("https must be 0(disabled) or 443(enabled)")
		os.Exit(-1)
	}
	if httpPort == 0 && httpsPort == 0 {
		fmt.Println("http or https are disabled")
		os.Exit(-1)
	}

	html := http.Dir("./html")
	fh := http.FileServer(html)
	http.Handle("/", fh)

	var protos []string
	if httpPort != 0 {
		protos = append(protos, fmt.Sprintf("http(:%v)", httpPort))
	}
	if httpsPort != 0 {
		protos = append(protos, fmt.Sprintf("https(:%v)", httpsPort))
	}
	fmt.Println(fmt.Sprintf("%v html root at %v", strings.Join(protos, ", "), string(html)))

	wg := sync.WaitGroup{}
	go func() {
		defer wg.Done()

		if httpPort == 0 {
			return
		}

		if err := http.ListenAndServe(fmt.Sprintf(":%v", httpPort), nil); err != nil {
			panic(err)
		}
	}()
	wg.Add(1)

	go func() {
		defer wg.Done()

		if httpsPort == 0 {
			return
		}

		var err error
		var m https.Manager
		if m, err = https.NewLetsencryptManager("", nil, "letsencrypt.cache"); err != nil {
			panic(err)
		}

		svr := &http.Server{
			Addr: fmt.Sprintf(":%v", httpsPort),
			TLSConfig: &tls.Config{
				GetCertificate: m.GetCertificate,
			},
		}

		if err := svr.ListenAndServeTLS("", ""); err != nil {
			panic(err)
		}
	}()

	wg.Wait()
}
<commit_msg>support domains<commit_after>/*
The MIT License (MIT)

Copyright (c) 2016 winlin

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

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

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

/*
 This the main entrance of https-proxy, proxy to api or other http server.
*/
package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/ossrs/go-oryx-lib/https"
	"net/http"
	"os"
	"strings"
	"sync"
)

func main() {
	var httpPort, httpsPort int
	var httpsDomains string
	flag.IntVar(&httpPort, "http", 80, "http listen at. 0 to disable http.")
	flag.IntVar(&httpsPort, "https", 443, "https listen at. 0 to disable https. 443 to serve. ")
	flag.StringVar(&httpsDomains, "domains", "", "the allow domains, empty to allow all. for example: ossrs.net,www.ossrs.net")
	flag.Parse()

	if httpsPort != 0 && httpsPort != 443 {
		fmt.Println("https must be 0(disabled) or 443(enabled)")
		os.Exit(-1)
	}
	if httpPort == 0 && httpsPort == 0 {
		fmt.Println("http or https are disabled")
		os.Exit(-1)
	}

	html := http.Dir("./html")
	fh := http.FileServer(html)
	http.Handle("/", fh)

	var protos []string
	if httpPort != 0 {
		protos = append(protos, fmt.Sprintf("http(:%v)", httpPort))
	}
	if httpsPort != 0 {
		s := httpsDomains
		if httpsDomains == "" {
			s = "all domains"
		}
		protos = append(protos, fmt.Sprintf("https(:%v, %v)", httpsPort, s))
	}
	fmt.Println(fmt.Sprintf("%v html root at %v", strings.Join(protos, ", "), string(html)))

	wg := sync.WaitGroup{}
	go func() {
		defer wg.Done()

		if httpPort == 0 {
			return
		}

		if err := http.ListenAndServe(fmt.Sprintf(":%v", httpPort), nil); err != nil {
			panic(err)
		}
	}()
	wg.Add(1)

	go func() {
		defer wg.Done()

		if httpsPort == 0 {
			return
		}

		var domains []string
		if httpsDomains != "" {
			domains = strings.Split(httpsDomains, ",")
		}

		var err error
		var m https.Manager
		if m, err = https.NewLetsencryptManager("", domains, "letsencrypt.cache"); err != nil {
			panic(err)
		}

		svr := &http.Server{
			Addr: fmt.Sprintf(":%v", httpsPort),
			TLSConfig: &tls.Config{
				GetCertificate: m.GetCertificate,
			},
		}

		if err := svr.ListenAndServeTLS("", ""); err != nil {
			panic(err)
		}
	}()

	wg.Wait()
}
<|endoftext|>


<commit_before>package eval

import (
	"fmt"

	"github.com/cirbo-lang/cirbo/cbty"
	"github.com/cirbo-lang/cirbo/source"
)

type StmtBlock struct {
	scope *Scope
	stmts []Stmt
}

// NewStmtBlock constructs and returns a new statement block containing the
// given statements, which are assumed to be populating the given scope.
//
// The caller must not read or write the given statements slice after it has
// been passed to NewStmtBlock. Ownership is transferred to the returned
// object and the slice's backing array may be modified in unspecified ways.
func MakeStmtBlock(scope *Scope, stmts []Stmt) (StmtBlock, source.Diags) {
	var diags source.Diags

	providers := make(map[*Symbol]Stmt, len(stmts))
	enables := make(map[Stmt][]Stmt, len(stmts)) // slice so that we preserve input ordering when ordering is ambiguous
	inDeg := make(map[Stmt]int, len(stmts))
	for _, stmt := range stmts {
		if sym := stmt.s.definedSymbol(); sym != nil {
			providers[sym] = stmt
		}
	}
	for _, stmt := range stmts {
		syms := stmt.s.requiredSymbols(scope)
		for sym := range syms {
			if provider, provided := providers[sym]; provided {
				enables[provider] = append(enables[provider], stmt)
				inDeg[stmt]++
			}
		}
	}

	// We place both "result" and "queue" at the head of our input array.
	// We know that the length of the queue and the length of the result
	// must sum up to less than or equal to the original list, so we can
	// safely use the original underlying array as storage for both. The
	// start of the queue will gradually move through the array just as
	// the result slice grows to include the elements it has vacated.
	result := stmts[0:0]
	queueStart := 0 // index into stmts underlying array
	queue := stmts[queueStart:queueStart]

	// Seed the queue with statements that have no dependencies
	for _, stmt := range stmts {
		if inDeg[stmt] == 0 {
			queue = append(queue, stmt)
		}
	}

	for len(queue) > 0 {
		stmt := queue[0]

		// Adjust the head of the queue to one element later in our array.
		queueStart++
		queue = stmts[queueStart : queueStart+(len(queue)-1)]

		// Adjust the result list to include the element that we just
		// removed from the queue.
		result = stmts[:len(result)+1]

		for _, enabled := range enables[stmt] {
			inDeg[enabled]--
			if inDeg[enabled] == 0 {
				queue = append(queue, enabled)
				delete(inDeg, enabled)
			}
		}
	}

	// When we reach this point, if there were no cycles then result already
	// equals stmts, but the list may have shrunk if there _were_ cycles and
	// so we need to do some adjusting.
	stmts = result

	if len(inDeg) > 0 {
		// Indicates that we have at least one cycle.
		// TODO: This error message isn't great; ideally we would provide
		// more context to help the user understand the reason for the
		// cycle, since it might be via multiple levels of indirection.
		ranges := make([]source.Range, 0, len(inDeg))
		for stmt := range inDeg {
			ranges = append(ranges, stmt.s.sourceRange())
		}

		if len(ranges) == 1 {
			diags = append(diags, source.Diag{
				Level:   source.Error,
				Summary: "Self-referential symbol definition",
				Detail:  "Definition statement depends (possibly indirectly) on its own result.",
				Ranges:  ranges,
			})
		} else {
			diags = append(diags, source.Diag{
				Level:   source.Error,
				Summary: "Self-referential symbol definitions",
				Detail:  "Definition statements depend (possibly indirectly) on their own results.",
				Ranges:  ranges,
			})
		}
	}

	// Don't permit any future modifications to the scope, since we're now
	// depending on its contents.
	scope.final = true

	return StmtBlock{
		scope: scope,
		stmts: stmts,
	}, diags
}

func (cb StmtBlock) RequiredSymbols(scope *Scope) SymbolSet {
	ret := NewSymbolSet()
	for _, stmt := range cb.stmts {
		reqd := stmt.RequiredSymbols(scope)
		for sym := range reqd {
			ret.Add(sym)
		}
	}
	return ret
}

func (sb StmtBlock) PackagesImported() []PackageRef {
	return sb.PackagesImportedAppend(nil)
}

func (sb StmtBlock) PackagesImportedAppend(ppaths []PackageRef) []PackageRef {
	for _, stmt := range sb.stmts {
		if imp, isImp := stmt.s.(*importStmt); isImp {
			ppaths = append(ppaths, PackageRef{
				Path:  imp.ppath,
				Range: imp.sourceRange(),
			})
		}
	}
	return ppaths
}

type PackageRef struct {
	Path  string
	Range source.Range
}

type StmtBlockAttr struct {
	Symbol   *Symbol
	Type     cbty.Type
	Required bool
	DefRange source.Range
}

// Attributes returns a description of the attributes defined by the block.
//
// When executing the block, values for some or all of these (depending on
// their Required status) should be provided in the StmtBlockExecute
// instance.
//
// The given context is used to resolve the type or default value expressions
// in the attribute statements. The given context must therefore be the same
// context that would ultimately be provided to Execute in the StmtBlockExecute
// object or else the result may be incorrect.
func (sb StmtBlock) Attributes(ctx *Context) (map[string]StmtBlockAttr, source.Diags) {
	var diags source.Diags
	ret := map[string]StmtBlockAttr{}
	for _, stmt := range sb.stmts {
		if attr, isAttr := stmt.s.(*attrStmt); isAttr {
			name := attr.sym.DeclaredName()
			def := StmtBlockAttr{
				Symbol: attr.sym,
			}

			switch {
			case attr.valueType != NilExpr:
				tyVal, exprDiags := attr.valueType.Value(ctx)
				diags = append(diags, exprDiags...)

				if tyVal.Type() != cbty.TypeType {
					diags = append(diags, source.Diag{
						Level:   source.Error,
						Summary: "Invalid attribute type",
						Detail:  fmt.Sprintf("Expected a type, but given a value of type %s. To assign a default value, use the '=' (equals) symbol.", tyVal.Type().Name()),
						Ranges:  attr.valueType.sourceRange().List(),
					})
					def.Type = cbty.PlaceholderVal.Type()
					def.Required = true
				} else {
					def.Type = tyVal.UnwrapType()
					def.Required = true
				}
			case attr.defValue != NilExpr:
				val, exprDiags := attr.defValue.Value(ctx)
				diags = append(diags, exprDiags...)

				def.Type = val.Type()
				def.Required = false
			default:
				// should never happen
				panic("attrStmt with neither value type nor default value")
			}

			ret[name] = def
		}
	}
	return ret, diags
}

// ImplicitExports returns a SymbolSet of the symbols defined in the block's
// scope that are eligible to be included in an implicit export object.
//
// This includes most definitions, but specifically excludes imports as they are
// assumed to be for internal use and could be requested directly by any caller.
//
// This method is intended for creating an implicit export object for a module,
// and so it will likely not produce a useful or sensible result for blocks
// created in other contexts.
func (sb StmtBlock) ImplicitExports() SymbolSet {
	var ret SymbolSet
	for _, stmt := range sb.stmts {
		if _, isImport := stmt.s.(*importStmt); isImport {
			continue
		}
		sym := stmt.s.definedSymbol()
		if sym != nil {
			if ret == nil {
				ret = SymbolSet{}
			}
			ret.Add(sym)
		}
	}
	return ret
}

func (sb StmtBlock) Execute(exec StmtBlockExecute) (*StmtBlockResult, source.Diags) {
	// Make a new child context to work in. (We get "exec" by value here, so
	// we can mutate it without upsetting the caller.)
	exec.Context = exec.Context.NewChild()

	result := StmtBlockResult{}
	var diags source.Diags

	result.Scope = sb.scope
	result.Context = exec.Context

	for _, stmt := range sb.stmts {
		stmtDiags := stmt.s.execute(&exec, &result)
		diags = append(diags, stmtDiags...)
	}

	result.Context.final = true // no more modifications allowed

	return &result, diags
}

type StmtBlockExecute struct {
	Context  *Context
	Packages map[string]cbty.Value
	Attrs    map[*Symbol]cbty.Value
}

type StmtBlockResult struct {
	// Context is the context that was created to evaluate the block.
	// It is provided only so that NewChild may be called on it for child
	// blocks; it should not be modified once returned.
	Context *Context

	// Scope is the scope that was created for the block during its
	// compilation. This object is shared between all executions of the same
	// block, and so should not be modified.
	Scope *Scope

	// ExportValue is the value exported by an "export" statement, if any.
	ExportValue cbty.Value
}
<commit_msg>eval: StmtBlock.AttributeNames<commit_after>package eval

import (
	"fmt"

	"github.com/cirbo-lang/cirbo/cbty"
	"github.com/cirbo-lang/cirbo/source"
)

type StmtBlock struct {
	scope *Scope
	stmts []Stmt
}

// NewStmtBlock constructs and returns a new statement block containing the
// given statements, which are assumed to be populating the given scope.
//
// The caller must not read or write the given statements slice after it has
// been passed to NewStmtBlock. Ownership is transferred to the returned
// object and the slice's backing array may be modified in unspecified ways.
func MakeStmtBlock(scope *Scope, stmts []Stmt) (StmtBlock, source.Diags) {
	var diags source.Diags

	providers := make(map[*Symbol]Stmt, len(stmts))
	enables := make(map[Stmt][]Stmt, len(stmts)) // slice so that we preserve input ordering when ordering is ambiguous
	inDeg := make(map[Stmt]int, len(stmts))
	for _, stmt := range stmts {
		if sym := stmt.s.definedSymbol(); sym != nil {
			providers[sym] = stmt
		}
	}
	for _, stmt := range stmts {
		syms := stmt.s.requiredSymbols(scope)
		for sym := range syms {
			if provider, provided := providers[sym]; provided {
				enables[provider] = append(enables[provider], stmt)
				inDeg[stmt]++
			}
		}
	}

	// We place both "result" and "queue" at the head of our input array.
	// We know that the length of the queue and the length of the result
	// must sum up to less than or equal to the original list, so we can
	// safely use the original underlying array as storage for both. The
	// start of the queue will gradually move through the array just as
	// the result slice grows to include the elements it has vacated.
	result := stmts[0:0]
	queueStart := 0 // index into stmts underlying array
	queue := stmts[queueStart:queueStart]

	// Seed the queue with statements that have no dependencies
	for _, stmt := range stmts {
		if inDeg[stmt] == 0 {
			queue = append(queue, stmt)
		}
	}

	for len(queue) > 0 {
		stmt := queue[0]

		// Adjust the head of the queue to one element later in our array.
		queueStart++
		queue = stmts[queueStart : queueStart+(len(queue)-1)]

		// Adjust the result list to include the element that we just
		// removed from the queue.
		result = stmts[:len(result)+1]

		for _, enabled := range enables[stmt] {
			inDeg[enabled]--
			if inDeg[enabled] == 0 {
				queue = append(queue, enabled)
				delete(inDeg, enabled)
			}
		}
	}

	// When we reach this point, if there were no cycles then result already
	// equals stmts, but the list may have shrunk if there _were_ cycles and
	// so we need to do some adjusting.
	stmts = result

	if len(inDeg) > 0 {
		// Indicates that we have at least one cycle.
		// TODO: This error message isn't great; ideally we would provide
		// more context to help the user understand the reason for the
		// cycle, since it might be via multiple levels of indirection.
		ranges := make([]source.Range, 0, len(inDeg))
		for stmt := range inDeg {
			ranges = append(ranges, stmt.s.sourceRange())
		}

		if len(ranges) == 1 {
			diags = append(diags, source.Diag{
				Level:   source.Error,
				Summary: "Self-referential symbol definition",
				Detail:  "Definition statement depends (possibly indirectly) on its own result.",
				Ranges:  ranges,
			})
		} else {
			diags = append(diags, source.Diag{
				Level:   source.Error,
				Summary: "Self-referential symbol definitions",
				Detail:  "Definition statements depend (possibly indirectly) on their own results.",
				Ranges:  ranges,
			})
		}
	}

	// Don't permit any future modifications to the scope, since we're now
	// depending on its contents.
	scope.final = true

	return StmtBlock{
		scope: scope,
		stmts: stmts,
	}, diags
}

func (cb StmtBlock) RequiredSymbols(scope *Scope) SymbolSet {
	ret := NewSymbolSet()
	for _, stmt := range cb.stmts {
		reqd := stmt.RequiredSymbols(scope)
		for sym := range reqd {
			ret.Add(sym)
		}
	}
	return ret
}

func (sb StmtBlock) PackagesImported() []PackageRef {
	return sb.PackagesImportedAppend(nil)
}

func (sb StmtBlock) PackagesImportedAppend(ppaths []PackageRef) []PackageRef {
	for _, stmt := range sb.stmts {
		if imp, isImp := stmt.s.(*importStmt); isImp {
			ppaths = append(ppaths, PackageRef{
				Path:  imp.ppath,
				Range: imp.sourceRange(),
			})
		}
	}
	return ppaths
}

type PackageRef struct {
	Path  string
	Range source.Range
}

type StmtBlockAttr struct {
	Symbol   *Symbol
	Type     cbty.Type
	Required bool
	DefRange source.Range
}

// AttributeNames returns a map of the attribute names required by the block,
// which can be resolved without needing a context. This can be used for
// early validation, though Attributes should be called with a context to
// get the full description of the attributes.
func (sb StmtBlock) AttributeNames() map[string]*Symbol {
	ret := map[string]*Symbol{}
	for _, stmt := range sb.stmts {
		if attr, isAttr := stmt.s.(*attrStmt); isAttr {
			ret[attr.sym.DeclaredName()] = attr.sym
		}
	}
	return ret
}

// Attributes returns a description of the attributes defined by the block.
//
// When executing the block, values for some or all of these (depending on
// their Required status) should be provided in the StmtBlockExecute
// instance.
//
// The given context is used to resolve the type or default value expressions
// in the attribute statements. The given context must therefore be the same
// context that would ultimately be provided to Execute in the StmtBlockExecute
// object or else the result may be incorrect.
func (sb StmtBlock) Attributes(ctx *Context) (map[string]StmtBlockAttr, source.Diags) {
	var diags source.Diags
	ret := map[string]StmtBlockAttr{}
	for _, stmt := range sb.stmts {
		if attr, isAttr := stmt.s.(*attrStmt); isAttr {
			name := attr.sym.DeclaredName()
			def := StmtBlockAttr{
				Symbol: attr.sym,
			}

			switch {
			case attr.valueType != NilExpr:
				tyVal, exprDiags := attr.valueType.Value(ctx)
				diags = append(diags, exprDiags...)

				if tyVal.Type() != cbty.TypeType {
					diags = append(diags, source.Diag{
						Level:   source.Error,
						Summary: "Invalid attribute type",
						Detail:  fmt.Sprintf("Expected a type, but given a value of type %s. To assign a default value, use the '=' (equals) symbol.", tyVal.Type().Name()),
						Ranges:  attr.valueType.sourceRange().List(),
					})
					def.Type = cbty.PlaceholderVal.Type()
					def.Required = true
				} else {
					def.Type = tyVal.UnwrapType()
					def.Required = true
				}
			case attr.defValue != NilExpr:
				val, exprDiags := attr.defValue.Value(ctx)
				diags = append(diags, exprDiags...)

				def.Type = val.Type()
				def.Required = false
			default:
				// should never happen
				panic("attrStmt with neither value type nor default value")
			}

			ret[name] = def
		}
	}
	return ret, diags
}

// ImplicitExports returns a SymbolSet of the symbols defined in the block's
// scope that are eligible to be included in an implicit export object.
//
// This includes most definitions, but specifically excludes imports as they are
// assumed to be for internal use and could be requested directly by any caller.
//
// This method is intended for creating an implicit export object for a module,
// and so it will likely not produce a useful or sensible result for blocks
// created in other contexts.
func (sb StmtBlock) ImplicitExports() SymbolSet {
	var ret SymbolSet
	for _, stmt := range sb.stmts {
		if _, isImport := stmt.s.(*importStmt); isImport {
			continue
		}
		sym := stmt.s.definedSymbol()
		if sym != nil {
			if ret == nil {
				ret = SymbolSet{}
			}
			ret.Add(sym)
		}
	}
	return ret
}

func (sb StmtBlock) Execute(exec StmtBlockExecute) (*StmtBlockResult, source.Diags) {
	// Make a new child context to work in. (We get "exec" by value here, so
	// we can mutate it without upsetting the caller.)
	exec.Context = exec.Context.NewChild()

	result := StmtBlockResult{}
	var diags source.Diags

	result.Scope = sb.scope
	result.Context = exec.Context

	for _, stmt := range sb.stmts {
		stmtDiags := stmt.s.execute(&exec, &result)
		diags = append(diags, stmtDiags...)
	}

	result.Context.final = true // no more modifications allowed

	return &result, diags
}

type StmtBlockExecute struct {
	Context  *Context
	Packages map[string]cbty.Value
	Attrs    map[*Symbol]cbty.Value
}

type StmtBlockResult struct {
	// Context is the context that was created to evaluate the block.
	// It is provided only so that NewChild may be called on it for child
	// blocks; it should not be modified once returned.
	Context *Context

	// Scope is the scope that was created for the block during its
	// compilation. This object is shared between all executions of the same
	// block, and so should not be modified.
	Scope *Scope

	// ExportValue is the value exported by an "export" statement, if any.
	ExportValue cbty.Value
}
<|endoftext|>