<commit_before>package mp

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

type MediaPlayer struct {
	player       Backend
	stateChange  chan StateChange
	volumeChange chan int

	// Two channels to handle passing of the play state without causing race
	// conditions.
	playstateChan     chan PlayState
	playstateChanPing chan bool

	streams map[string]string // map of YouTube ID to stream gotten from youtube-dl
}

func New(stateChange chan StateChange, volumeChange chan int) *MediaPlayer {
	p := MediaPlayer{}

	p.playstateChan = make(chan PlayState)
	p.playstateChanPing = make(chan bool)
	p.stateChange = stateChange
	p.volumeChange = volumeChange
	p.streams = make(map[string]string)

	p.player = &VLC{}
	playerEventChan := p.player.initialize()

	go p.run(playerEventChan)

	return &p
}

func (p *MediaPlayer) Quit() {
	go func() {
		// TODO: fix race conditions (the player might not be fully initialized yet)
		p.player.quit()
	}()

	// TODO: fix race conditions
	close(p.stateChange)
	close(p.volumeChange)
}

func (p *MediaPlayer) GetPosition() time.Duration {
	ps := p.GetPlaystate()

	return p.getPosition(ps)
}

func (p *MediaPlayer) getPosition(ps *PlayState) time.Duration {
	var position time.Duration
	if ps.State == STATE_STOPPED {
		position = 0
	} else if ps.State == STATE_BUFFERING {
		position = ps.bufferingPosition
	} else {
		position = p.player.getPosition()
	}

	if position < 0 {
		panic("got position < 0")
	}

	return position
}

// GetPlaystate returns the play state synchronously (may block)
func (p *MediaPlayer) GetPlaystate() *PlayState {
	p.playstateChanPing <- false
	ps := <-p.playstateChan
	return &ps
}

// changePlaystate changes the play state inside a callback
// The *PlayState argument in the callback is the PlayState that can be
// changed, but it can only be accessed inside the callback (outside of
// it, race conditions can occur).
func (p *MediaPlayer) changePlaystate(callback func(*PlayState)) {
	p.playstateChanPing <- true
	ps := <-p.playstateChan
	callback(&ps)
	p.playstateChan <- ps
}

// SetPlaystate changes the play state to the specified arguments
// This function doesn't block, but changes may not be immediately applied.
func (p *MediaPlayer) SetPlaystate(playlist []string, index int, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State == STATE_BUFFERING && ps.bufferingPosition == position && playlist[index] == ps.Playlist[ps.Index] {
			// just in case something else has changed, update the playlist
			p.updatePlaylist(ps, playlist)
			return
		}
		ps.Playlist = playlist
		ps.Index = index

		if len(ps.Playlist) > 0 {
			p.startPlaying(ps, position)
		} else {
			p.Stop()
		}
	})
}

func (p *MediaPlayer) startPlaying(ps *PlayState, position time.Duration) {
	p.setPlayState(ps, STATE_BUFFERING, position)

	videoId := ps.Playlist[ps.Index]

	// do not use the playstate inside this goroutine to prevent race conditions
	go func() {
		streamUrl := p.getYouTubeStream(videoId)
		p.player.play(streamUrl, position)
	}()
}

func (p *MediaPlayer) getYouTubeStream(videoId string) string {
	if stream, ok := p.streams[videoId]; ok {
		return stream
	}

	youtubeUrl := "https://www.youtube.com/watch?v=" + videoId

	fmt.Println("Fetching YouTube stream...", youtubeUrl)
	// First (mkv-container) audio only, then video with audio bitrate 100+
	// (where video has the lowest possible quality), then slightly lower
	// quality audio.
	// We do this because for some reason DASH aac audio (in the MP4 container)
	// doesn't support seeking in any of the tested players (mpv using
	// libavformat, and vlc, gstreamer and mplayer2 using their own demuxers).
	// But the MKV container seems to have much better support.
	// See:
	//   https://github.com/mpv-player/mpv/issues/579
	//   https://trac.ffmpeg.org/ticket/3842
	cmd := exec.Command("youtube-dl", "-g", "-f", "171/172/43/22/18", youtubeUrl)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	fmt.Println("Got stream.")
	if err != nil {
		panic(err)
	}
	stream := strings.TrimSpace(string(output))
	p.streams[videoId] = stream
	return stream
}

// setPlayState sets updates the PlayState and sends events.
// position may be -1: in that case it will be updated.
func (p *MediaPlayer) setPlayState(ps *PlayState, state State, position time.Duration) {
	if state == ps.State {
		fmt.Printf("WARNING: state %d did not change\n", state)
	}
	ps.State = state

	if state == STATE_BUFFERING {
		ps.bufferingPosition = position
	} else {
		ps.bufferingPosition = -1
	}

	if position == -1 {
		position = p.getPosition(ps)
	}

	go func() {
		p.stateChange <- StateChange{state, position}
	}()
}

func (p *MediaPlayer) UpdatePlaylist(playlist []string) {
	go p.changePlaystate(func(ps *PlayState) {
		p.updatePlaylist(ps, playlist)
	})
}

func (p *MediaPlayer) updatePlaylist(ps *PlayState, playlist []string) {
	if len(ps.Playlist) == 0 {

		if ps.State == STATE_PLAYING {
			// just to be sure
			panic("empty playlist while playing")
		}
		ps.Playlist = playlist

		if ps.Index >= len(playlist) {
			// this appears to be the normal behavior of YouTube
			ps.Index = len(playlist) - 1
		}

		if ps.State == STATE_STOPPED {
			p.startPlaying(ps, 0)
		}

	} else {
		videoId := ps.Playlist[ps.Index]
		ps.Playlist = playlist
		p.setPlaylistIndex(ps, videoId)
	}
}

func (p *MediaPlayer) SetVideo(videoId string, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		p.setPlaylistIndex(ps, videoId)
		p.startPlaying(ps, position)
	})
}

func (p *MediaPlayer) setPlaylistIndex(ps *PlayState, videoId string) {
	newIndex := -1
	for i, v := range ps.Playlist {
		if v == videoId {
			if newIndex >= 0 {
				fmt.Fprintln(os.Stderr, "WARNING: videoId exists twice in playlist")
				break
			}
			newIndex = i
			// no 'break' so duplicate video entries can be checked
		}
	}

	if newIndex == -1 {
		// don't know how to proceed
		panic("current video does not exist in new playlist")
	}

	ps.Index = newIndex
}

// Pause pauses the currently playing video
func (p *MediaPlayer) Pause() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PLAYING {
			fmt.Printf("Warning: pause while in state %d\n", ps.State)
		}

		p.player.pause()
	})
}

// Resume resumes playback when it was paused
func (p *MediaPlayer) Resume() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: resume while in state %d\n", ps.State)
		}

		p.player.resume()
	})
}

// Seek jumps to the specified position
func (p *MediaPlayer) Seek(position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: state is not paused while seeking (state: %d)\n", ps.State)
		}

		p.player.setPosition(position)
	})
}

// SetVolume sets the volume of the player to the specified value (0-100).
func (p *MediaPlayer) SetVolume(volume int) {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Volume = volume
		p.player.setVolume(ps.Volume)
	})
}

// ChangeVolume increases or decreases the volume by the specified delta.
// It returns a channel with the result, that can be ignored.
func (p *MediaPlayer) ChangeVolume(delta int) chan int {
	// the buffer makes sure a send will succeed even if the channel is never
	// read: this way the channel can be garbage collected and won't leave a
	// garbage goroutine if it is never read.
	c := make(chan int, 1)
	go p.changePlaystate(func(ps *PlayState) {

		ps.Volume += delta
		// pressing 'volume up' or 'volume down' keeps sending volume
		// increase/decrease messages. Keep the volume within range 0-100.
		if ps.Volume < 0 {
			ps.Volume = 0
		}
		if ps.Volume > 100 {
			ps.Volume = 100
		}

		p.player.setVolume(ps.Volume)

		c <- ps.Volume
	})
	return c
}

func (p *MediaPlayer) Stop() {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = []string{}
		// Do not set ps.Index to 0, it may be needed for UpdatePlaylist:
		// Stop is called before UpdatePlaylist when removing the currently
		// playing video from the playlist.
		p.player.stop()
	})
}

func (p *MediaPlayer) run(playerEventChan chan State) {
	ps := PlayState{}

	ps.Volume = INITIAL_VOLUME

	for {
		select {
		case replace := <-p.playstateChanPing:
			p.playstateChan <- ps
			if replace {
				ps = <-p.playstateChan
			}

		case event, ok := <-playerEventChan:
			if !ok {
				// player has quit, and closed channel
				return
			}

			switch event {
			case STATE_PLAYING:
				p.setPlayState(&ps, STATE_PLAYING, -1)

			case STATE_PAUSED:
				p.setPlayState(&ps, STATE_PAUSED, -1)

			case STATE_STOPPED:
				if ps.State == STATE_BUFFERING {
					// Especially VLC may keep sending 'stopped' events
					// while the next track is already buffering.
					// Ignore those events.
					continue
				}
				if ps.Index+1 < len(ps.Playlist) {
					// there are more videos, play the next
					ps.Index++
					p.startPlaying(&ps, 0)
				} else {
					// signal that the video has stopped playing
					// this resets the position but keeps the playlist
					p.setPlayState(&ps, STATE_STOPPED, 0)
				}
			}
		}
	}
}
<commit_msg>Don't warn: it is quite normal for events to be sent multiple times<commit_after>package mp

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

type MediaPlayer struct {
	player       Backend
	stateChange  chan StateChange
	volumeChange chan int

	// Two channels to handle passing of the play state without causing race
	// conditions.
	playstateChan     chan PlayState
	playstateChanPing chan bool

	streams map[string]string // map of YouTube ID to stream gotten from youtube-dl
}

func New(stateChange chan StateChange, volumeChange chan int) *MediaPlayer {
	p := MediaPlayer{}

	p.playstateChan = make(chan PlayState)
	p.playstateChanPing = make(chan bool)
	p.stateChange = stateChange
	p.volumeChange = volumeChange
	p.streams = make(map[string]string)

	p.player = &VLC{}
	playerEventChan := p.player.initialize()

	go p.run(playerEventChan)

	return &p
}

func (p *MediaPlayer) Quit() {
	go func() {
		// TODO: fix race conditions (the player might not be fully initialized yet)
		p.player.quit()
	}()

	// TODO: fix race conditions
	close(p.stateChange)
	close(p.volumeChange)
}

func (p *MediaPlayer) GetPosition() time.Duration {
	ps := p.GetPlaystate()

	return p.getPosition(ps)
}

func (p *MediaPlayer) getPosition(ps *PlayState) time.Duration {
	var position time.Duration
	if ps.State == STATE_STOPPED {
		position = 0
	} else if ps.State == STATE_BUFFERING {
		position = ps.bufferingPosition
	} else {
		position = p.player.getPosition()
	}

	if position < 0 {
		panic("got position < 0")
	}

	return position
}

// GetPlaystate returns the play state synchronously (may block)
func (p *MediaPlayer) GetPlaystate() *PlayState {
	p.playstateChanPing <- false
	ps := <-p.playstateChan
	return &ps
}

// changePlaystate changes the play state inside a callback
// The *PlayState argument in the callback is the PlayState that can be
// changed, but it can only be accessed inside the callback (outside of
// it, race conditions can occur).
func (p *MediaPlayer) changePlaystate(callback func(*PlayState)) {
	p.playstateChanPing <- true
	ps := <-p.playstateChan
	callback(&ps)
	p.playstateChan <- ps
}

// SetPlaystate changes the play state to the specified arguments
// This function doesn't block, but changes may not be immediately applied.
func (p *MediaPlayer) SetPlaystate(playlist []string, index int, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State == STATE_BUFFERING && ps.bufferingPosition == position && playlist[index] == ps.Playlist[ps.Index] {
			// just in case something else has changed, update the playlist
			p.updatePlaylist(ps, playlist)
			return
		}
		ps.Playlist = playlist
		ps.Index = index

		if len(ps.Playlist) > 0 {
			p.startPlaying(ps, position)
		} else {
			p.Stop()
		}
	})
}

func (p *MediaPlayer) startPlaying(ps *PlayState, position time.Duration) {
	p.setPlayState(ps, STATE_BUFFERING, position)

	videoId := ps.Playlist[ps.Index]

	// do not use the playstate inside this goroutine to prevent race conditions
	go func() {
		streamUrl := p.getYouTubeStream(videoId)
		p.player.play(streamUrl, position)
	}()
}

func (p *MediaPlayer) getYouTubeStream(videoId string) string {
	if stream, ok := p.streams[videoId]; ok {
		return stream
	}

	youtubeUrl := "https://www.youtube.com/watch?v=" + videoId

	fmt.Println("Fetching YouTube stream...", youtubeUrl)
	// First (mkv-container) audio only, then video with audio bitrate 100+
	// (where video has the lowest possible quality), then slightly lower
	// quality audio.
	// We do this because for some reason DASH aac audio (in the MP4 container)
	// doesn't support seeking in any of the tested players (mpv using
	// libavformat, and vlc, gstreamer and mplayer2 using their own demuxers).
	// But the MKV container seems to have much better support.
	// See:
	//   https://github.com/mpv-player/mpv/issues/579
	//   https://trac.ffmpeg.org/ticket/3842
	cmd := exec.Command("youtube-dl", "-g", "-f", "171/172/43/22/18", youtubeUrl)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	fmt.Println("Got stream.")
	if err != nil {
		panic(err)
	}
	stream := strings.TrimSpace(string(output))
	p.streams[videoId] = stream
	return stream
}

// setPlayState sets updates the PlayState and sends events.
// position may be -1: in that case it will be updated.
func (p *MediaPlayer) setPlayState(ps *PlayState, state State, position time.Duration) {
	ps.State = state

	if state == STATE_BUFFERING {
		ps.bufferingPosition = position
	} else {
		ps.bufferingPosition = -1
	}

	if position == -1 {
		position = p.getPosition(ps)
	}

	go func() {
		p.stateChange <- StateChange{state, position}
	}()
}

func (p *MediaPlayer) UpdatePlaylist(playlist []string) {
	go p.changePlaystate(func(ps *PlayState) {
		p.updatePlaylist(ps, playlist)
	})
}

func (p *MediaPlayer) updatePlaylist(ps *PlayState, playlist []string) {
	if len(ps.Playlist) == 0 {

		if ps.State == STATE_PLAYING {
			// just to be sure
			panic("empty playlist while playing")
		}
		ps.Playlist = playlist

		if ps.Index >= len(playlist) {
			// this appears to be the normal behavior of YouTube
			ps.Index = len(playlist) - 1
		}

		if ps.State == STATE_STOPPED {
			p.startPlaying(ps, 0)
		}

	} else {
		videoId := ps.Playlist[ps.Index]
		ps.Playlist = playlist
		p.setPlaylistIndex(ps, videoId)
	}
}

func (p *MediaPlayer) SetVideo(videoId string, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		p.setPlaylistIndex(ps, videoId)
		p.startPlaying(ps, position)
	})
}

func (p *MediaPlayer) setPlaylistIndex(ps *PlayState, videoId string) {
	newIndex := -1
	for i, v := range ps.Playlist {
		if v == videoId {
			if newIndex >= 0 {
				fmt.Fprintln(os.Stderr, "WARNING: videoId exists twice in playlist")
				break
			}
			newIndex = i
			// no 'break' so duplicate video entries can be checked
		}
	}

	if newIndex == -1 {
		// don't know how to proceed
		panic("current video does not exist in new playlist")
	}

	ps.Index = newIndex
}

// Pause pauses the currently playing video
func (p *MediaPlayer) Pause() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PLAYING {
			fmt.Printf("Warning: pause while in state %d\n", ps.State)
		}

		p.player.pause()
	})
}

// Resume resumes playback when it was paused
func (p *MediaPlayer) Resume() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: resume while in state %d\n", ps.State)
		}

		p.player.resume()
	})
}

// Seek jumps to the specified position
func (p *MediaPlayer) Seek(position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			fmt.Printf("Warning: state is not paused while seeking (state: %d)\n", ps.State)
		}

		p.player.setPosition(position)
	})
}

// SetVolume sets the volume of the player to the specified value (0-100).
func (p *MediaPlayer) SetVolume(volume int) {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Volume = volume
		p.player.setVolume(ps.Volume)
	})
}

// ChangeVolume increases or decreases the volume by the specified delta.
// It returns a channel with the result, that can be ignored.
func (p *MediaPlayer) ChangeVolume(delta int) chan int {
	// the buffer makes sure a send will succeed even if the channel is never
	// read: this way the channel can be garbage collected and won't leave a
	// garbage goroutine if it is never read.
	c := make(chan int, 1)
	go p.changePlaystate(func(ps *PlayState) {

		ps.Volume += delta
		// pressing 'volume up' or 'volume down' keeps sending volume
		// increase/decrease messages. Keep the volume within range 0-100.
		if ps.Volume < 0 {
			ps.Volume = 0
		}
		if ps.Volume > 100 {
			ps.Volume = 100
		}

		p.player.setVolume(ps.Volume)

		c <- ps.Volume
	})
	return c
}

func (p *MediaPlayer) Stop() {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = []string{}
		// Do not set ps.Index to 0, it may be needed for UpdatePlaylist:
		// Stop is called before UpdatePlaylist when removing the currently
		// playing video from the playlist.
		p.player.stop()
	})
}

func (p *MediaPlayer) run(playerEventChan chan State) {
	ps := PlayState{}

	ps.Volume = INITIAL_VOLUME

	for {
		select {
		case replace := <-p.playstateChanPing:
			p.playstateChan <- ps
			if replace {
				ps = <-p.playstateChan
			}

		case event, ok := <-playerEventChan:
			if !ok {
				// player has quit, and closed channel
				return
			}

			switch event {
			case STATE_PLAYING:
				p.setPlayState(&ps, STATE_PLAYING, -1)

			case STATE_PAUSED:
				p.setPlayState(&ps, STATE_PAUSED, -1)

			case STATE_STOPPED:
				if ps.State == STATE_BUFFERING {
					// Especially VLC may keep sending 'stopped' events
					// while the next track is already buffering.
					// Ignore those events.
					continue
				}
				if ps.Index+1 < len(ps.Playlist) {
					// there are more videos, play the next
					ps.Index++
					p.startPlaying(&ps, 0)
				} else {
					// signal that the video has stopped playing
					// this resets the position but keeps the playlist
					p.setPlayState(&ps, STATE_STOPPED, 0)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package mp

import (
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

type MediaPlayer struct {
	player       Backend
	stateChange  chan StateChange
	volumeChange chan int

	// Two channels to handle passing of the play state without causing race
	// conditions.
	playstateChan     chan PlayState
	playstateChanPing chan bool

	streams map[string]string // map of YouTube ID to stream gotten from youtube-dl
}

func New(stateChange chan StateChange, volumeChange chan int) *MediaPlayer {
	p := MediaPlayer{}

	p.playstateChan = make(chan PlayState)
	p.playstateChanPing = make(chan bool)
	p.stateChange = stateChange
	p.volumeChange = volumeChange
	p.streams = make(map[string]string)

	p.player = &VLC{}
	playerEventChan := p.player.initialize()

	go p.run(playerEventChan)

	return &p
}

func (p *MediaPlayer) Quit() {
	go func() {
		// TODO: fix race conditions (the player might not be fully initialized yet)
		p.player.quit()
	}()

	// TODO: fix race conditions
	close(p.stateChange)
	close(p.volumeChange)
}

func (p *MediaPlayer) GetPosition() time.Duration {
	ps := p.GetPlaystate()

	return p.getPosition(ps)
}

func (p *MediaPlayer) getPosition(ps *PlayState) time.Duration {
	var position time.Duration
	if ps.State == STATE_STOPPED {
		position = 0
	} else if ps.State == STATE_BUFFERING {
		position = ps.bufferingPosition
	} else {
		position = p.player.getPosition()
	}

	if position < 0 {
		panic("got position < 0")
	}

	return position
}

// GetPlaystate returns the play state synchronously (may block)
func (p *MediaPlayer) GetPlaystate() *PlayState {
	p.playstateChanPing <- false
	ps := <-p.playstateChan
	return &ps
}

// changePlaystate changes the play state inside a callback
// The *PlayState argument in the callback is the PlayState that can be
// changed, but it can only be accessed inside the callback (outside of
// it, race conditions can occur).
func (p *MediaPlayer) changePlaystate(callback func(*PlayState)) {
	p.playstateChanPing <- true
	ps := <-p.playstateChan
	callback(&ps)
	p.playstateChan <- ps
}

// SetPlaystate changes the play state to the specified arguments
// This function doesn't block, but changes may not be immediately applied.
func (p *MediaPlayer) SetPlaystate(playlist []string, index int, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State == STATE_BUFFERING && ps.bufferingPosition == position && playlist[index] == ps.Playlist[ps.Index] {
			// just in case something else has changed, update the playlist
			p.updatePlaylist(ps, playlist)
			return
		}
		ps.Playlist = playlist
		ps.Index = index

		if len(ps.Playlist) > 0 {
			p.startPlaying(ps, position)
		} else {
			p.Stop()
		}
	})
}

func (p *MediaPlayer) startPlaying(ps *PlayState, position time.Duration) {
	if ps.State == STATE_PLAYING {
		// Pause the currently playing track.
		// This has multiple benefits:
		//  *  When pressing 'play', the user probably expects the next video to
		//     be played immediately, or if that is not possible, expects the
		//     current video to stop playing.
		//  *  On very slow systems, like the Raspberry Pi, downloading the
		//     stream URL for the next video doesn't interrupt the currently
		//     playing video.
		p.player.pause()
	}
	p.setPlayState(ps, STATE_BUFFERING, position)

	videoId := ps.Playlist[ps.Index]

	// do not use the playstate inside this goroutine to prevent race conditions
	go func() {
		streamUrl := p.getYouTubeStream(videoId)
		p.player.play(streamUrl, position)
	}()
}

func (p *MediaPlayer) getYouTubeStream(videoId string) string {
	if stream, ok := p.streams[videoId]; ok {
		return stream
	}

	youtubeUrl := "https://www.youtube.com/watch?v=" + videoId

	log.Println("Fetching YouTube stream...", youtubeUrl)
	// First (mkv-container) audio only, then video with audio bitrate 100+
	// (where video has the lowest possible quality), then slightly lower
	// quality audio.
	// We do this because for some reason DASH aac audio (in the MP4 container)
	// doesn't support seeking in any of the tested players (mpv using
	// libavformat, and vlc, gstreamer and mplayer2 using their own demuxers).
	// But the MKV container seems to have much better support.
	// See:
	//   https://github.com/mpv-player/mpv/issues/579
	//   https://trac.ffmpeg.org/ticket/3842
	cmd := exec.Command("youtube-dl", "-g", "-f", "171/172/43/22/18", youtubeUrl)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	log.Println("Got stream.")
	if err != nil {
		panic(err)
	}
	stream := strings.TrimSpace(string(output))
	p.streams[videoId] = stream
	return stream
}

// setPlayState sets updates the PlayState and sends events.
// position may be -1: in that case it will be updated.
func (p *MediaPlayer) setPlayState(ps *PlayState, state State, position time.Duration) {
	ps.State = state

	if state == STATE_BUFFERING {
		ps.bufferingPosition = position
	} else {
		ps.bufferingPosition = -1
	}

	if position == -1 {
		position = p.getPosition(ps)
	}

	go func() {
		p.stateChange <- StateChange{state, position}
	}()
}

func (p *MediaPlayer) UpdatePlaylist(playlist []string) {
	go p.changePlaystate(func(ps *PlayState) {
		p.updatePlaylist(ps, playlist)
	})
}

func (p *MediaPlayer) updatePlaylist(ps *PlayState, playlist []string) {
	if len(ps.Playlist) == 0 {

		if ps.State == STATE_PLAYING {
			// just to be sure
			panic("empty playlist while playing")
		}
		ps.Playlist = playlist

		if ps.Index >= len(playlist) {
			// this appears to be the normal behavior of YouTube
			ps.Index = len(playlist) - 1
		}

		if ps.State == STATE_STOPPED {
			p.startPlaying(ps, 0)
		}

	} else {
		videoId := ps.Playlist[ps.Index]
		ps.Playlist = playlist
		p.setPlaylistIndex(ps, videoId)
	}
}

func (p *MediaPlayer) SetVideo(videoId string, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		p.setPlaylistIndex(ps, videoId)
		p.startPlaying(ps, position)
	})
}

func (p *MediaPlayer) setPlaylistIndex(ps *PlayState, videoId string) {
	newIndex := -1
	for i, v := range ps.Playlist {
		if v == videoId {
			if newIndex >= 0 {
				log.Println("WARNING: videoId exists twice in playlist")
				break
			}
			newIndex = i
			// no 'break' so duplicate video entries can be checked
		}
	}

	if newIndex == -1 {
		// don't know how to proceed
		panic("current video does not exist in new playlist")
	}

	ps.Index = newIndex
}

// Pause pauses the currently playing video
func (p *MediaPlayer) Pause() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PLAYING {
			log.Printf("Warning: pause while in state %d\n", ps.State)
		}

		p.player.pause()
	})
}

// Play resumes playback when it was paused
func (p *MediaPlayer) Play() {
	go p.changePlaystate(func(ps *PlayState) {
		switch ps.State {
		case STATE_PAUSED:
			p.player.resume()
		case STATE_STOPPED:
			// Restart from the beginning.
			p.startPlaying(ps, 0)
		default:
			log.Printf("Warning: resume while in state %d\n", ps.State)
		}
	})
}

// Seek jumps to the specified position
func (p *MediaPlayer) Seek(position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			log.Printf("Warning: state is not paused while seeking (state: %d)\n", ps.State)
		}

		p.player.setPosition(position)
	})
}

// SetVolume sets the volume of the player to the specified value (0-100).
func (p *MediaPlayer) SetVolume(volume int) {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Volume = volume
		p.player.setVolume(ps.Volume)
	})
}

// ChangeVolume increases or decreases the volume by the specified delta.
// It returns a channel with the result, that can be ignored.
func (p *MediaPlayer) ChangeVolume(delta int) chan int {
	// the buffer makes sure a send will succeed even if the channel is never
	// read: this way the channel can be garbage collected and won't leave a
	// garbage goroutine if it is never read.
	c := make(chan int, 1)
	go p.changePlaystate(func(ps *PlayState) {

		ps.Volume += delta
		// pressing 'volume up' or 'volume down' keeps sending volume
		// increase/decrease messages. Keep the volume within range 0-100.
		if ps.Volume < 0 {
			ps.Volume = 0
		}
		if ps.Volume > 100 {
			ps.Volume = 100
		}

		p.player.setVolume(ps.Volume)

		c <- ps.Volume
	})
	return c
}

func (p *MediaPlayer) Stop() {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = []string{}
		// Do not set ps.Index to 0, it may be needed for UpdatePlaylist:
		// Stop is called before UpdatePlaylist when removing the currently
		// playing video from the playlist.
		p.player.stop()
	})
}

func (p *MediaPlayer) run(playerEventChan chan State) {
	ps := PlayState{}

	ps.Volume = INITIAL_VOLUME

	for {
		select {
		case replace := <-p.playstateChanPing:
			p.playstateChan <- ps
			if replace {
				ps = <-p.playstateChan
			}

		case event, ok := <-playerEventChan:
			if !ok {
				// player has quit, and closed channel
				return
			}

			switch event {
			case STATE_PLAYING:
				p.setPlayState(&ps, STATE_PLAYING, -1)

			case STATE_PAUSED:
				p.setPlayState(&ps, STATE_PAUSED, -1)

			case STATE_STOPPED:
				if ps.State == STATE_BUFFERING {
					// Especially VLC may keep sending 'stopped' events
					// while the next track is already buffering.
					// Ignore those events.
					fmt.Println("WARNING: 'stopped' event while buffering")
					continue
				}
				if ps.Index+1 < len(ps.Playlist) {
					// there are more videos, play the next
					ps.Index++
					p.startPlaying(&ps, 0)
				} else {
					// signal that the video has stopped playing
					// this resets the position but keeps the playlist
					p.setPlayState(&ps, STATE_STOPPED, 0)
				}
			}
		}
	}
}
<commit_msg>This isn't a race condition AFAICS<commit_after>package mp

import (
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

type MediaPlayer struct {
	player       Backend
	stateChange  chan StateChange
	volumeChange chan int

	// Two channels to handle passing of the play state without causing race
	// conditions.
	playstateChan     chan PlayState
	playstateChanPing chan bool

	streams map[string]string // map of YouTube ID to stream gotten from youtube-dl
}

func New(stateChange chan StateChange, volumeChange chan int) *MediaPlayer {
	p := MediaPlayer{}

	p.playstateChan = make(chan PlayState)
	p.playstateChanPing = make(chan bool)
	p.stateChange = stateChange
	p.volumeChange = volumeChange
	p.streams = make(map[string]string)

	p.player = &VLC{}
	playerEventChan := p.player.initialize()

	go p.run(playerEventChan)

	return &p
}

func (p *MediaPlayer) Quit() {
	go func() {
		// TODO: fix race conditions (the player might not be fully initialized yet)
		p.player.quit()
	}()

	close(p.stateChange)
	close(p.volumeChange)
}

func (p *MediaPlayer) GetPosition() time.Duration {
	ps := p.GetPlaystate()

	return p.getPosition(ps)
}

func (p *MediaPlayer) getPosition(ps *PlayState) time.Duration {
	var position time.Duration
	if ps.State == STATE_STOPPED {
		position = 0
	} else if ps.State == STATE_BUFFERING {
		position = ps.bufferingPosition
	} else {
		position = p.player.getPosition()
	}

	if position < 0 {
		panic("got position < 0")
	}

	return position
}

// GetPlaystate returns the play state synchronously (may block)
func (p *MediaPlayer) GetPlaystate() *PlayState {
	p.playstateChanPing <- false
	ps := <-p.playstateChan
	return &ps
}

// changePlaystate changes the play state inside a callback
// The *PlayState argument in the callback is the PlayState that can be
// changed, but it can only be accessed inside the callback (outside of
// it, race conditions can occur).
func (p *MediaPlayer) changePlaystate(callback func(*PlayState)) {
	p.playstateChanPing <- true
	ps := <-p.playstateChan
	callback(&ps)
	p.playstateChan <- ps
}

// SetPlaystate changes the play state to the specified arguments
// This function doesn't block, but changes may not be immediately applied.
func (p *MediaPlayer) SetPlaystate(playlist []string, index int, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State == STATE_BUFFERING && ps.bufferingPosition == position && playlist[index] == ps.Playlist[ps.Index] {
			// just in case something else has changed, update the playlist
			p.updatePlaylist(ps, playlist)
			return
		}
		ps.Playlist = playlist
		ps.Index = index

		if len(ps.Playlist) > 0 {
			p.startPlaying(ps, position)
		} else {
			p.Stop()
		}
	})
}

func (p *MediaPlayer) startPlaying(ps *PlayState, position time.Duration) {
	if ps.State == STATE_PLAYING {
		// Pause the currently playing track.
		// This has multiple benefits:
		//  *  When pressing 'play', the user probably expects the next video to
		//     be played immediately, or if that is not possible, expects the
		//     current video to stop playing.
		//  *  On very slow systems, like the Raspberry Pi, downloading the
		//     stream URL for the next video doesn't interrupt the currently
		//     playing video.
		p.player.pause()
	}
	p.setPlayState(ps, STATE_BUFFERING, position)

	videoId := ps.Playlist[ps.Index]

	// do not use the playstate inside this goroutine to prevent race conditions
	go func() {
		streamUrl := p.getYouTubeStream(videoId)
		p.player.play(streamUrl, position)
	}()
}

func (p *MediaPlayer) getYouTubeStream(videoId string) string {
	if stream, ok := p.streams[videoId]; ok {
		return stream
	}

	youtubeUrl := "https://www.youtube.com/watch?v=" + videoId

	log.Println("Fetching YouTube stream...", youtubeUrl)
	// First (mkv-container) audio only, then video with audio bitrate 100+
	// (where video has the lowest possible quality), then slightly lower
	// quality audio.
	// We do this because for some reason DASH aac audio (in the MP4 container)
	// doesn't support seeking in any of the tested players (mpv using
	// libavformat, and vlc, gstreamer and mplayer2 using their own demuxers).
	// But the MKV container seems to have much better support.
	// See:
	//   https://github.com/mpv-player/mpv/issues/579
	//   https://trac.ffmpeg.org/ticket/3842
	cmd := exec.Command("youtube-dl", "-g", "-f", "171/172/43/22/18", youtubeUrl)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	log.Println("Got stream.")
	if err != nil {
		panic(err)
	}
	stream := strings.TrimSpace(string(output))
	p.streams[videoId] = stream
	return stream
}

// setPlayState sets updates the PlayState and sends events.
// position may be -1: in that case it will be updated.
func (p *MediaPlayer) setPlayState(ps *PlayState, state State, position time.Duration) {
	ps.State = state

	if state == STATE_BUFFERING {
		ps.bufferingPosition = position
	} else {
		ps.bufferingPosition = -1
	}

	if position == -1 {
		position = p.getPosition(ps)
	}

	go func() {
		p.stateChange <- StateChange{state, position}
	}()
}

func (p *MediaPlayer) UpdatePlaylist(playlist []string) {
	go p.changePlaystate(func(ps *PlayState) {
		p.updatePlaylist(ps, playlist)
	})
}

func (p *MediaPlayer) updatePlaylist(ps *PlayState, playlist []string) {
	if len(ps.Playlist) == 0 {

		if ps.State == STATE_PLAYING {
			// just to be sure
			panic("empty playlist while playing")
		}
		ps.Playlist = playlist

		if ps.Index >= len(playlist) {
			// this appears to be the normal behavior of YouTube
			ps.Index = len(playlist) - 1
		}

		if ps.State == STATE_STOPPED {
			p.startPlaying(ps, 0)
		}

	} else {
		videoId := ps.Playlist[ps.Index]
		ps.Playlist = playlist
		p.setPlaylistIndex(ps, videoId)
	}
}

func (p *MediaPlayer) SetVideo(videoId string, position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		p.setPlaylistIndex(ps, videoId)
		p.startPlaying(ps, position)
	})
}

func (p *MediaPlayer) setPlaylistIndex(ps *PlayState, videoId string) {
	newIndex := -1
	for i, v := range ps.Playlist {
		if v == videoId {
			if newIndex >= 0 {
				log.Println("WARNING: videoId exists twice in playlist")
				break
			}
			newIndex = i
			// no 'break' so duplicate video entries can be checked
		}
	}

	if newIndex == -1 {
		// don't know how to proceed
		panic("current video does not exist in new playlist")
	}

	ps.Index = newIndex
}

// Pause pauses the currently playing video
func (p *MediaPlayer) Pause() {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PLAYING {
			log.Printf("Warning: pause while in state %d\n", ps.State)
		}

		p.player.pause()
	})
}

// Play resumes playback when it was paused
func (p *MediaPlayer) Play() {
	go p.changePlaystate(func(ps *PlayState) {
		switch ps.State {
		case STATE_PAUSED:
			p.player.resume()
		case STATE_STOPPED:
			// Restart from the beginning.
			p.startPlaying(ps, 0)
		default:
			log.Printf("Warning: resume while in state %d\n", ps.State)
		}
	})
}

// Seek jumps to the specified position
func (p *MediaPlayer) Seek(position time.Duration) {
	go p.changePlaystate(func(ps *PlayState) {
		if ps.State != STATE_PAUSED {
			log.Printf("Warning: state is not paused while seeking (state: %d)\n", ps.State)
		}

		p.player.setPosition(position)
	})
}

// SetVolume sets the volume of the player to the specified value (0-100).
func (p *MediaPlayer) SetVolume(volume int) {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Volume = volume
		p.player.setVolume(ps.Volume)
	})
}

// ChangeVolume increases or decreases the volume by the specified delta.
// It returns a channel with the result, that can be ignored.
func (p *MediaPlayer) ChangeVolume(delta int) chan int {
	// the buffer makes sure a send will succeed even if the channel is never
	// read: this way the channel can be garbage collected and won't leave a
	// garbage goroutine if it is never read.
	c := make(chan int, 1)
	go p.changePlaystate(func(ps *PlayState) {

		ps.Volume += delta
		// pressing 'volume up' or 'volume down' keeps sending volume
		// increase/decrease messages. Keep the volume within range 0-100.
		if ps.Volume < 0 {
			ps.Volume = 0
		}
		if ps.Volume > 100 {
			ps.Volume = 100
		}

		p.player.setVolume(ps.Volume)

		c <- ps.Volume
	})
	return c
}

func (p *MediaPlayer) Stop() {
	go p.changePlaystate(func(ps *PlayState) {
		ps.Playlist = []string{}
		// Do not set ps.Index to 0, it may be needed for UpdatePlaylist:
		// Stop is called before UpdatePlaylist when removing the currently
		// playing video from the playlist.
		p.player.stop()
	})
}

func (p *MediaPlayer) run(playerEventChan chan State) {
	ps := PlayState{}

	ps.Volume = INITIAL_VOLUME

	for {
		select {
		case replace := <-p.playstateChanPing:
			p.playstateChan <- ps
			if replace {
				ps = <-p.playstateChan
			}

		case event, ok := <-playerEventChan:
			if !ok {
				// player has quit, and closed channel
				return
			}

			switch event {
			case STATE_PLAYING:
				p.setPlayState(&ps, STATE_PLAYING, -1)

			case STATE_PAUSED:
				p.setPlayState(&ps, STATE_PAUSED, -1)

			case STATE_STOPPED:
				if ps.State == STATE_BUFFERING {
					// Especially VLC may keep sending 'stopped' events
					// while the next track is already buffering.
					// Ignore those events.
					fmt.Println("WARNING: 'stopped' event while buffering")
					continue
				}
				if ps.Index+1 < len(ps.Playlist) {
					// there are more videos, play the next
					ps.Index++
					p.startPlaying(&ps, 0)
				} else {
					// signal that the video has stopped playing
					// this resets the position but keeps the playlist
					p.setPlayState(&ps, STATE_STOPPED, 0)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2016

package instana

import (
	"errors"
	"net/http"
	"strings"

	"github.com/instana/go-sensor/w3ctrace"
	ot "github.com/opentracing/opentracing-go"
)

// Instana header constants
const (
	// FieldT Trace ID header
	FieldT = "x-instana-t"
	// FieldS Span ID header
	FieldS = "x-instana-s"
	// FieldL Level header
	FieldL = "x-instana-l"
	// FieldB OT Baggage header
	FieldB = "x-instana-b-"
	// FieldSynthetic if set to 1, marks the call as synthetic, e.g.
	// a healthcheck request
	FieldSynthetic = "x-instana-synthetic"
)

func injectTraceContext(sc SpanContext, opaqueCarrier interface{}) error {
	roCarrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	// Handle pre-existing case-sensitive keys
	exstfieldT := FieldT
	exstfieldS := FieldS
	exstfieldL := FieldL
	exstfieldB := FieldB

	roCarrier.ForeachKey(func(k, v string) error {
		switch strings.ToLower(k) {
		case FieldT:
			exstfieldT = k
		case FieldS:
			exstfieldS = k
		case FieldL:
			exstfieldL = k
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				exstfieldB = string([]rune(k)[:len(FieldB)])
			}
		}
		return nil
	})

	carrier, ok := opaqueCarrier.(ot.TextMapWriter)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		// Even though the godoc claims that the key passed to (*http.Header).Set()
		// is case-insensitive, it actually normalizes it using textproto.CanonicalMIMEHeaderKey()
		// before populating the value. As a result headers with non-canonical will not be
		// overwritten with a new value. This is only the case if header names were set while
		// initializing the http.Header instance, i.e.
		//     h := http.Headers{"X-InStAnA-T": {"abc123"}}
		// and does not apply to a common case when requests are being created using http.NewRequest()
		// or http.ReadRequest() that call (*http.Header).Set() to set header values.
		h := http.Header(c)
		delete(h, exstfieldT)
		delete(h, exstfieldS)
		delete(h, exstfieldL)

		for key := range h {
			if strings.HasPrefix(strings.ToLower(key), FieldB) {
				delete(h, key)
			}
		}

		addW3CTraceContext(h, sc)
		addEUMHeaders(h, sc)
	}

	if !sc.Suppressed {
		carrier.Set(exstfieldT, FormatID(sc.TraceID))
		carrier.Set(exstfieldS, FormatID(sc.SpanID))
	} else {
		// remove trace context keys from the carrier
		switch c := opaqueCarrier.(type) {
		case ot.HTTPHeadersCarrier:
			h := http.Header(c)
			h.Del(exstfieldT)
			h.Del(exstfieldS)
		case ot.TextMapCarrier:
			delete(c, exstfieldT)
			delete(c, exstfieldS)
		case interface{ RemoveAll() }:
			// in case carrier has the RemoveAll() method that wipes all trace
			// headers, for example the instasarama.ProducerMessagCarrier, we
			// use it to remove the context of a suppressed trace
			c.RemoveAll()
		}
	}

	carrier.Set(exstfieldL, formatLevel(sc))

	for k, v := range sc.Baggage {
		carrier.Set(exstfieldB+k, v)
	}

	return nil
}

func extractTraceContext(opaqueCarrier interface{}) (SpanContext, error) {
	spanContext := SpanContext{
		Baggage: make(map[string]string),
	}

	carrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return spanContext, ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		pickupW3CTraceContext(http.Header(c), &spanContext)
	}

	err := carrier.ForeachKey(func(k, v string) error {
		var err error

		switch strings.ToLower(k) {
		case FieldT:
			spanContext.TraceIDHi, spanContext.TraceID, err = ParseLongID(v)
			if err != nil {
				return ot.ErrSpanContextCorrupted
			}
		case FieldS:
			spanContext.SpanID, err = ParseID(v)
			if err != nil {
				return ot.ErrSpanContextCorrupted
			}
		case FieldL:
			suppressed, corrData, err := parseLevel(v)
			if err != nil {
				sensor.logger.Info("failed to parse ", k, ": ", err, " (", v, ")")
				// use defaults
				suppressed, corrData = false, EUMCorrelationData{}
			}

			spanContext.Suppressed = suppressed
			if !spanContext.Suppressed {
				spanContext.Correlation = corrData
			}
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				// preserve original case of the baggage key
				spanContext.Baggage[k[len(FieldB):]] = v
			}
		}

		return nil
	})
	if err != nil {
		return spanContext, err
	}

	// reset the trace if a correlation ID has been provided
	if spanContext.Correlation.ID != "" {
		spanContext.TraceIDHi, spanContext.TraceID, spanContext.SpanID = 0, 0, 0

		return spanContext, nil
	}

	if spanContext.IsZero() {
		return spanContext, ot.ErrSpanContextNotFound
	}

	if !spanContext.Suppressed && (spanContext.SpanID == 0 != (spanContext.TraceIDHi == 0 && spanContext.TraceID == 0)) {
		return spanContext, ot.ErrSpanContextCorrupted
	}

	return spanContext, nil
}

func addW3CTraceContext(h http.Header, sc SpanContext) {
	traceID, spanID := FormatLongID(sc.TraceIDHi, sc.TraceID), FormatID(sc.SpanID)
	trCtx := sc.W3CContext

	// check for an existing w3c trace
	if trCtx.IsZero() {
		// initiate trace if none
		trCtx = w3ctrace.New(w3ctrace.Parent{
			Version:  w3ctrace.Version_Max,
			TraceID:  traceID,
			ParentID: spanID,
			Flags: w3ctrace.Flags{
				Sampled: !sc.Suppressed,
			},
		})
	}

	// update the traceparent parent ID
	p := trCtx.Parent()
	p.ParentID = spanID
	// sync the traceparent `sampled` flags with the X-Instana-L value
	p.Flags.Sampled = !sc.Suppressed

	trCtx.RawParent = p.String()

	// participate in w3c trace context if tracing is enabled
	if !sc.Suppressed {
		// propagate truncated trace ID downstream
		trCtx.RawState = trCtx.State().Add(w3ctrace.VendorInstana, FormatID(sc.TraceID)+";"+spanID).String()
	}

	w3ctrace.Inject(trCtx, h)
}

func pickupW3CTraceContext(h http.Header, sc *SpanContext) {
	trCtx, err := w3ctrace.Extract(h)
	if err != nil {
		return
	}
	sc.W3CContext = trCtx
}

func addEUMHeaders(h http.Header, sc SpanContext) {
	// Preserve original Server-Timing header values by combining them into a comma-separated list
	st := append(h["Server-Timing"], "intid;desc="+FormatID(sc.TraceID))
	h.Set("Server-Timing", strings.Join(st, ", "))
}

var errMalformedHeader = errors.New("malformed header value")

func parseLevel(s string) (bool, EUMCorrelationData, error) {
	const (
		levelState uint8 = iota
		partSeparatorState
		correlationPartState
		correlationTypeState
		correlationIDState
		finalState
	)

	if s == "" {
		return false, EUMCorrelationData{}, nil
	}

	var (
		typeInd                 int
		state                   uint8
		level, corrType, corrID string
	)
PARSE:
	for ptr := 0; state != finalState && ptr < len(s); ptr++ {
		switch state {
		case levelState: // looking for 0 or 1
			level = s[ptr : ptr+1]

			if level != "0" && level != "1" {
				break PARSE
			}

			if ptr == len(s)-1 { // no correlation ID provided
				state = finalState
			} else {
				state = partSeparatorState
			}
		case partSeparatorState: // skip OWS while looking for ','
			switch s[ptr] {
			case ' ', '\t': // advance
			case ',':
				state = correlationPartState
			default:
				break PARSE
			}
		case correlationPartState: // skip OWS while searching for 'correlationType=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // advance
			case strings.HasPrefix(s[ptr:], "correlationType="):
				ptr += 15 // advance to the end of prefix
				typeInd = ptr + 1
				state = correlationTypeState
			default:
				break PARSE
			}
		case correlationTypeState: // skip OWS while looking for ';'
			switch s[ptr] {
			case ' ', '\t': // possibly trailing OWS, advance
			case ';':
				state = correlationIDState
			default:
				corrType = s[typeInd : ptr+1]
			}
		case correlationIDState: //  skip OWS while searching for 'correlationId=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // leading OWS, advance
			case strings.HasPrefix(s[ptr:], "correlationId="):
				ptr += 14
				corrID = s[ptr:]
				state = finalState
			default:
				break PARSE
			}
		default:
			break PARSE
		}
	}

	if state != finalState {
		return false, EUMCorrelationData{}, errMalformedHeader
	}

	return level == "0", EUMCorrelationData{Type: corrType, ID: corrID}, nil
}

func formatLevel(sc SpanContext) string {
	if sc.Suppressed {
		return "0"
	}

	return "1"
}
<commit_msg>return from tracerS.Extract spanContext, when Instana headers were corrupted but w3 context was found<commit_after>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2016

package instana

import (
	"errors"
	"net/http"
	"strings"

	"github.com/instana/go-sensor/w3ctrace"
	ot "github.com/opentracing/opentracing-go"
)

// Instana header constants
const (
	// FieldT Trace ID header
	FieldT = "x-instana-t"
	// FieldS Span ID header
	FieldS = "x-instana-s"
	// FieldL Level header
	FieldL = "x-instana-l"
	// FieldB OT Baggage header
	FieldB = "x-instana-b-"
	// FieldSynthetic if set to 1, marks the call as synthetic, e.g.
	// a healthcheck request
	FieldSynthetic = "x-instana-synthetic"
)

func injectTraceContext(sc SpanContext, opaqueCarrier interface{}) error {
	roCarrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	// Handle pre-existing case-sensitive keys
	exstfieldT := FieldT
	exstfieldS := FieldS
	exstfieldL := FieldL
	exstfieldB := FieldB

	roCarrier.ForeachKey(func(k, v string) error {
		switch strings.ToLower(k) {
		case FieldT:
			exstfieldT = k
		case FieldS:
			exstfieldS = k
		case FieldL:
			exstfieldL = k
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				exstfieldB = string([]rune(k)[:len(FieldB)])
			}
		}
		return nil
	})

	carrier, ok := opaqueCarrier.(ot.TextMapWriter)
	if !ok {
		return ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		// Even though the godoc claims that the key passed to (*http.Header).Set()
		// is case-insensitive, it actually normalizes it using textproto.CanonicalMIMEHeaderKey()
		// before populating the value. As a result headers with non-canonical will not be
		// overwritten with a new value. This is only the case if header names were set while
		// initializing the http.Header instance, i.e.
		//     h := http.Headers{"X-InStAnA-T": {"abc123"}}
		// and does not apply to a common case when requests are being created using http.NewRequest()
		// or http.ReadRequest() that call (*http.Header).Set() to set header values.
		h := http.Header(c)
		delete(h, exstfieldT)
		delete(h, exstfieldS)
		delete(h, exstfieldL)

		for key := range h {
			if strings.HasPrefix(strings.ToLower(key), FieldB) {
				delete(h, key)
			}
		}

		addW3CTraceContext(h, sc)
		addEUMHeaders(h, sc)
	}

	if !sc.Suppressed {
		carrier.Set(exstfieldT, FormatID(sc.TraceID))
		carrier.Set(exstfieldS, FormatID(sc.SpanID))
	} else {
		// remove trace context keys from the carrier
		switch c := opaqueCarrier.(type) {
		case ot.HTTPHeadersCarrier:
			h := http.Header(c)
			h.Del(exstfieldT)
			h.Del(exstfieldS)
		case ot.TextMapCarrier:
			delete(c, exstfieldT)
			delete(c, exstfieldS)
		case interface{ RemoveAll() }:
			// in case carrier has the RemoveAll() method that wipes all trace
			// headers, for example the instasarama.ProducerMessagCarrier, we
			// use it to remove the context of a suppressed trace
			c.RemoveAll()
		}
	}

	carrier.Set(exstfieldL, formatLevel(sc))

	for k, v := range sc.Baggage {
		carrier.Set(exstfieldB+k, v)
	}

	return nil
}

func extractTraceContext(opaqueCarrier interface{}) (SpanContext, error) {
	spanContext := SpanContext{
		Baggage: make(map[string]string),
	}

	carrier, ok := opaqueCarrier.(ot.TextMapReader)
	if !ok {
		return spanContext, ot.ErrInvalidCarrier
	}

	if c, ok := opaqueCarrier.(ot.HTTPHeadersCarrier); ok {
		pickupW3CTraceContext(http.Header(c), &spanContext)
	}

	err := carrier.ForeachKey(func(k, v string) error {
		var err error

		switch strings.ToLower(k) {
		case FieldT:
			spanContext.TraceIDHi, spanContext.TraceID, err = ParseLongID(v)
			if err != nil {
				return ot.ErrSpanContextCorrupted
			}
		case FieldS:
			spanContext.SpanID, err = ParseID(v)
			if err != nil {
				return ot.ErrSpanContextCorrupted
			}
		case FieldL:
			suppressed, corrData, err := parseLevel(v)
			if err != nil {
				sensor.logger.Info("failed to parse ", k, ": ", err, " (", v, ")")
				// use defaults
				suppressed, corrData = false, EUMCorrelationData{}
			}

			spanContext.Suppressed = suppressed
			if !spanContext.Suppressed {
				spanContext.Correlation = corrData
			}
		default:
			if strings.HasPrefix(strings.ToLower(k), FieldB) {
				// preserve original case of the baggage key
				spanContext.Baggage[k[len(FieldB):]] = v
			}
		}

		return nil
	})
	if err != nil {
		return spanContext, err
	}

	// reset the trace if a correlation ID has been provided
	if spanContext.Correlation.ID != "" {
		spanContext.TraceIDHi, spanContext.TraceID, spanContext.SpanID = 0, 0, 0

		return spanContext, nil
	}

	if spanContext.IsZero() {
		return spanContext, ot.ErrSpanContextNotFound
	}

	// when the context is not suppressed
	// and instana headers either can not be parsed or present partially
	// and w3 context is not there
	if !spanContext.Suppressed &&
		(spanContext.SpanID == 0 != (spanContext.TraceIDHi == 0 && spanContext.TraceID == 0)) &&
		spanContext.W3CContext.IsZero() {

		return spanContext, ot.ErrSpanContextCorrupted
	}

	return spanContext, nil
}

func addW3CTraceContext(h http.Header, sc SpanContext) {
	traceID, spanID := FormatLongID(sc.TraceIDHi, sc.TraceID), FormatID(sc.SpanID)
	trCtx := sc.W3CContext

	// check for an existing w3c trace
	if trCtx.IsZero() {
		// initiate trace if none
		trCtx = w3ctrace.New(w3ctrace.Parent{
			Version:  w3ctrace.Version_Max,
			TraceID:  traceID,
			ParentID: spanID,
			Flags: w3ctrace.Flags{
				Sampled: !sc.Suppressed,
			},
		})
	}

	// update the traceparent parent ID
	p := trCtx.Parent()
	p.ParentID = spanID
	// sync the traceparent `sampled` flags with the X-Instana-L value
	p.Flags.Sampled = !sc.Suppressed

	trCtx.RawParent = p.String()

	// participate in w3c trace context if tracing is enabled
	if !sc.Suppressed {
		// propagate truncated trace ID downstream
		trCtx.RawState = trCtx.State().Add(w3ctrace.VendorInstana, FormatID(sc.TraceID)+";"+spanID).String()
	}

	w3ctrace.Inject(trCtx, h)
}

func pickupW3CTraceContext(h http.Header, sc *SpanContext) {
	trCtx, err := w3ctrace.Extract(h)
	if err != nil {
		return
	}
	sc.W3CContext = trCtx
}

func addEUMHeaders(h http.Header, sc SpanContext) {
	// Preserve original Server-Timing header values by combining them into a comma-separated list
	st := append(h["Server-Timing"], "intid;desc="+FormatID(sc.TraceID))
	h.Set("Server-Timing", strings.Join(st, ", "))
}

var errMalformedHeader = errors.New("malformed header value")

func parseLevel(s string) (bool, EUMCorrelationData, error) {
	const (
		levelState uint8 = iota
		partSeparatorState
		correlationPartState
		correlationTypeState
		correlationIDState
		finalState
	)

	if s == "" {
		return false, EUMCorrelationData{}, nil
	}

	var (
		typeInd                 int
		state                   uint8
		level, corrType, corrID string
	)
PARSE:
	for ptr := 0; state != finalState && ptr < len(s); ptr++ {
		switch state {
		case levelState: // looking for 0 or 1
			level = s[ptr : ptr+1]

			if level != "0" && level != "1" {
				break PARSE
			}

			if ptr == len(s)-1 { // no correlation ID provided
				state = finalState
			} else {
				state = partSeparatorState
			}
		case partSeparatorState: // skip OWS while looking for ','
			switch s[ptr] {
			case ' ', '\t': // advance
			case ',':
				state = correlationPartState
			default:
				break PARSE
			}
		case correlationPartState: // skip OWS while searching for 'correlationType=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // advance
			case strings.HasPrefix(s[ptr:], "correlationType="):
				ptr += 15 // advance to the end of prefix
				typeInd = ptr + 1
				state = correlationTypeState
			default:
				break PARSE
			}
		case correlationTypeState: // skip OWS while looking for ';'
			switch s[ptr] {
			case ' ', '\t': // possibly trailing OWS, advance
			case ';':
				state = correlationIDState
			default:
				corrType = s[typeInd : ptr+1]
			}
		case correlationIDState: //  skip OWS while searching for 'correlationId=' prefix
			switch {
			case s[ptr] == ' ' || s[ptr] == '\t': // leading OWS, advance
			case strings.HasPrefix(s[ptr:], "correlationId="):
				ptr += 14
				corrID = s[ptr:]
				state = finalState
			default:
				break PARSE
			}
		default:
			break PARSE
		}
	}

	if state != finalState {
		return false, EUMCorrelationData{}, errMalformedHeader
	}

	return level == "0", EUMCorrelationData{Type: corrType, ID: corrID}, nil
}

func formatLevel(sc SpanContext) string {
	if sc.Suppressed {
		return "0"
	}

	return "1"
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/base64"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"sync/atomic"
)

type Proxy struct {
	transcoders map[string]Transcoder
	ml          *mitmListener
	ReadCount   uint64
	WriteCount  uint64
	user        string
	pass        string
}

type Transcoder interface {
	Transcode(*ResponseWriter, *ResponseReader, http.Header) error
}

func New() *Proxy {
	p := &Proxy{
		transcoders: make(map[string]Transcoder),
		ml:          nil,
	}
	return p
}

func (p *Proxy) EnableMitm(ca, key string) error {
	cf, err := newCertFaker(ca, key)
	if err != nil {
		return err
	}
	p.ml = newMitmListener(cf)
	go http.Serve(p.ml, p)
	return nil
}

func (p *Proxy) SetAuthentication(user, pass string) {
	p.user = user
	p.pass = pass
}

func (p *Proxy) AddTranscoder(contentType string, transcoder Transcoder) {
	p.transcoders[contentType] = transcoder
}

func (p *Proxy) Start(host string) error {
	return http.ListenAndServe(host, p)
}

func (p *Proxy) StartTLS(host, cert, key string) error {
	return http.ListenAndServeTLS(host, cert, key, p)
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("serving request: %s", r.URL)
	if err := p.handle(w, r); err != nil {
		log.Printf("%s while serving request: %s", err, r.URL)
	}
}

func (p *Proxy) checkHttpBasicAuth(auth string) bool {
	prefix := "Basic "
	if !strings.HasPrefix(auth, prefix) {
		return false
	}
	decoded, err := base64.StdEncoding.DecodeString(auth[len(prefix):])
	if err != nil {
		return false
	}
	values := strings.SplitN(string(decoded), ":", 2)
	if len(values) != 2 || values[0] != p.user || values[1] != p.pass {
		return false
	}
	return true
}

func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error {
	if r.Method == "CONNECT" {
		return p.handleConnect(w, r)
	}

	// TODO: only HTTPS?
	if p.user != "" {
		if !p.checkHttpBasicAuth(r.Header.Get("Proxy-Authorization")) {
			w.Header().Set("WWW-Authenticate", "Basic realm=\"Compy\"")
			w.WriteHeader(http.StatusProxyAuthRequired)
			return nil
		}
	}

	resp, err := forward(r)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return fmt.Errorf("error forwarding request: %s", err)
	}
	defer resp.Body.Close()
	rw := newResponseWriter(w)
	rr := newResponseReader(resp)
	err = p.proxyResponse(rw, rr, r.Header)
	read := rr.counter.Count()
	written := rw.rw.Count()
	log.Printf("transcoded: %d -> %d (%3.1f%%)", read, written, float64(written)/float64(read)*100)
	atomic.AddUint64(&p.ReadCount, read)
	atomic.AddUint64(&p.WriteCount, written)
	return err
}

func forward(r *http.Request) (*http.Response, error) {
	if r.URL.Scheme == "" {
		if r.TLS != nil && r.TLS.ServerName == r.Host {
			r.URL.Scheme = "https"
		} else {
			r.URL.Scheme = "http"
		}
	}
	if r.URL.Host == "" {
		r.URL.Host = r.Host
	}
	r.RequestURI = ""
	return http.DefaultTransport.RoundTrip(r)
}

func (p *Proxy) proxyResponse(w *ResponseWriter, r *ResponseReader, headers http.Header) error {
	w.takeHeaders(r)
	transcoder, found := p.transcoders[r.ContentType()]
	if !found {
		return w.ReadFrom(r)
	}
	w.setChunked()
	if err := transcoder.Transcode(w, r, headers); err != nil {
		return fmt.Errorf("transcoding error: %s", err)
	}
	return nil
}

func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) error {
	if p.ml == nil {
		return fmt.Errorf("CONNECT received but mitm is not enabled")
	}
	w.WriteHeader(http.StatusOK)
	var conn net.Conn
	if h, ok := w.(http.Hijacker); ok {
		conn, _, _ = h.Hijack()
	} else {
		fw := w.(FlushWriter)
		fw.Flush()
		mconn := newMitmConn(fw, r.Body, r.RemoteAddr)
		conn = mconn
		defer func() {
			<-mconn.closed
		}()
	}
	sconn, err := p.ml.Serve(conn, r.Host)
	if err != nil {
		conn.Close()
		return err
	}
	sconn.Close() // TODO: reuse this connection for https requests
	return nil
}
<commit_msg>Check HTTP auth before allowing CONNECT<commit_after>package proxy

import (
	"encoding/base64"
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"
	"sync/atomic"
)

type Proxy struct {
	transcoders map[string]Transcoder
	ml          *mitmListener
	ReadCount   uint64
	WriteCount  uint64
	user        string
	pass        string
}

type Transcoder interface {
	Transcode(*ResponseWriter, *ResponseReader, http.Header) error
}

func New() *Proxy {
	p := &Proxy{
		transcoders: make(map[string]Transcoder),
		ml:          nil,
	}
	return p
}

func (p *Proxy) EnableMitm(ca, key string) error {
	cf, err := newCertFaker(ca, key)
	if err != nil {
		return err
	}
	p.ml = newMitmListener(cf)
	go http.Serve(p.ml, p)
	return nil
}

func (p *Proxy) SetAuthentication(user, pass string) {
	p.user = user
	p.pass = pass
}

func (p *Proxy) AddTranscoder(contentType string, transcoder Transcoder) {
	p.transcoders[contentType] = transcoder
}

func (p *Proxy) Start(host string) error {
	return http.ListenAndServe(host, p)
}

func (p *Proxy) StartTLS(host, cert, key string) error {
	return http.ListenAndServeTLS(host, cert, key, p)
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Printf("serving request: %s", r.URL)
	if err := p.handle(w, r); err != nil {
		log.Printf("%s while serving request: %s", err, r.URL)
	}
}

func (p *Proxy) checkHttpBasicAuth(auth string) bool {
	prefix := "Basic "
	if !strings.HasPrefix(auth, prefix) {
		return false
	}
	decoded, err := base64.StdEncoding.DecodeString(auth[len(prefix):])
	if err != nil {
		return false
	}
	values := strings.SplitN(string(decoded), ":", 2)
	if len(values) != 2 || values[0] != p.user || values[1] != p.pass {
		return false
	}
	return true
}

func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error {
	// TODO: only HTTPS?
	if p.user != "" {
		if !p.checkHttpBasicAuth(r.Header.Get("Proxy-Authorization")) {
			w.Header().Set("WWW-Authenticate", "Basic realm=\"Compy\"")
			w.WriteHeader(http.StatusProxyAuthRequired)
			return nil
		}
	}

	if r.Method == "CONNECT" {
		return p.handleConnect(w, r)
	}

	resp, err := forward(r)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return fmt.Errorf("error forwarding request: %s", err)
	}
	defer resp.Body.Close()
	rw := newResponseWriter(w)
	rr := newResponseReader(resp)
	err = p.proxyResponse(rw, rr, r.Header)
	read := rr.counter.Count()
	written := rw.rw.Count()
	log.Printf("transcoded: %d -> %d (%3.1f%%)", read, written, float64(written)/float64(read)*100)
	atomic.AddUint64(&p.ReadCount, read)
	atomic.AddUint64(&p.WriteCount, written)
	return err
}

func forward(r *http.Request) (*http.Response, error) {
	if r.URL.Scheme == "" {
		if r.TLS != nil && r.TLS.ServerName == r.Host {
			r.URL.Scheme = "https"
		} else {
			r.URL.Scheme = "http"
		}
	}
	if r.URL.Host == "" {
		r.URL.Host = r.Host
	}
	r.RequestURI = ""
	return http.DefaultTransport.RoundTrip(r)
}

func (p *Proxy) proxyResponse(w *ResponseWriter, r *ResponseReader, headers http.Header) error {
	w.takeHeaders(r)
	transcoder, found := p.transcoders[r.ContentType()]
	if !found {
		return w.ReadFrom(r)
	}
	w.setChunked()
	if err := transcoder.Transcode(w, r, headers); err != nil {
		return fmt.Errorf("transcoding error: %s", err)
	}
	return nil
}

func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) error {
	if p.ml == nil {
		return fmt.Errorf("CONNECT received but mitm is not enabled")
	}
	w.WriteHeader(http.StatusOK)
	var conn net.Conn
	if h, ok := w.(http.Hijacker); ok {
		conn, _, _ = h.Hijack()
	} else {
		fw := w.(FlushWriter)
		fw.Flush()
		mconn := newMitmConn(fw, r.Body, r.RemoteAddr)
		conn = mconn
		defer func() {
			<-mconn.closed
		}()
	}
	sconn, err := p.ml.Serve(conn, r.Host)
	if err != nil {
		conn.Close()
		return err
	}
	sconn.Close() // TODO: reuse this connection for https requests
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/workspaces"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"

	"github.com/terraform-providers/terraform-provider-aws/aws/internal/keyvaluetags"
	"github.com/terraform-providers/terraform-provider-aws/aws/internal/service/workspaces/waiter"
)

func resourceAwsWorkspacesWorkspace() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsWorkspacesWorkspaceCreate,
		Read:   resourceAwsWorkspacesWorkspaceRead,
		Update: resourceAwsWorkspacesWorkspaceUpdate,
		Delete: resourceAwsWorkspacesWorkspaceDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},
		Schema: map[string]*schema.Schema{
			"bundle_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"directory_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"ip_address": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"computer_name": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"state": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"root_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
				Default:  false,
			},
			"user_name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"user_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
				Default:  false,
			},
			"volume_encryption_key": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},
			"workspace_properties": {
				Type:     schema.TypeList,
				MaxItems: 1,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"compute_type_name": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  workspaces.ComputeValue,
							ValidateFunc: validation.StringInSlice([]string{
								workspaces.ComputeValue,
								workspaces.ComputeStandard,
								workspaces.ComputePerformance,
								workspaces.ComputePower,
								workspaces.ComputePowerpro,
								workspaces.ComputeGraphics,
								workspaces.ComputeGraphicspro,
							}, false),
						},
						"root_volume_size_gib": {
							Type:     schema.TypeInt,
							Optional: true,
							Default:  80,
							ValidateFunc: validation.Any(
								validation.IntInSlice([]int{80}),
								validation.IntBetween(175, 2000),
							),
						},
						"running_mode": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  workspaces.RunningModeAlwaysOn,
							ValidateFunc: validation.StringInSlice([]string{
								workspaces.RunningModeAlwaysOn,
								workspaces.RunningModeAutoStop,
							}, false),
						},
						"running_mode_auto_stop_timeout_in_minutes": {
							Type:     schema.TypeInt,
							Optional: true,
							Computed: true,
							ValidateFunc: func(v interface{}, k string) (ws []string, errors []error) {
								val := v.(int)
								if val%60 != 0 {
									errors = append(errors, fmt.Errorf(
										"%q should be configured in 60-minute intervals, got: %d", k, val))
								}
								return
							},
						},
						"user_volume_size_gib": {
							Type:     schema.TypeInt,
							Optional: true,
							Default:  10,
							ValidateFunc: validation.Any(
								validation.IntInSlice([]int{10, 50}),
								validation.IntBetween(100, 2000),
							),
						},
					},
				},
			},
			"tags": tagsSchema(),
		},
	}
}

func resourceAwsWorkspacesWorkspaceCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn

	tags := keyvaluetags.New(d.Get("tags").(map[string]interface{})).IgnoreAws().WorkspacesTags()

	input := &workspaces.WorkspaceRequest{
		BundleId:                    aws.String(d.Get("bundle_id").(string)),
		DirectoryId:                 aws.String(d.Get("directory_id").(string)),
		UserName:                    aws.String(d.Get("user_name").(string)),
		RootVolumeEncryptionEnabled: aws.Bool(d.Get("root_volume_encryption_enabled").(bool)),
		UserVolumeEncryptionEnabled: aws.Bool(d.Get("user_volume_encryption_enabled").(bool)),
		Tags:                        tags,
	}

	if v, ok := d.GetOk("volume_encryption_key"); ok {
		input.VolumeEncryptionKey = aws.String(v.(string))
	}

	input.WorkspaceProperties = expandWorkspaceProperties(d.Get("workspace_properties").([]interface{}))

	log.Printf("[DEBUG] Creating workspace...\n%#v\n", *input)
	resp, err := conn.CreateWorkspaces(&workspaces.CreateWorkspacesInput{
		Workspaces: []*workspaces.WorkspaceRequest{input},
	})
	if err != nil {
		return err
	}

	wsFail := resp.FailedRequests
	if len(wsFail) > 0 {
		return fmt.Errorf("workspace creation failed: %s", *wsFail[0].ErrorMessage)
	}

	workspaceID := aws.StringValue(resp.PendingRequests[0].WorkspaceId)

	log.Printf("[DEBUG] Waiting for workspace %q to be available...", workspaceID)
	_, err = waiter.WorkspaceAvailable(conn, workspaceID)
	if err != nil {
		return fmt.Errorf("workspace %q is not available: %s", workspaceID, err)
	}

	d.SetId(workspaceID)
	log.Printf("[DEBUG] Workspace %q is available", workspaceID)

	return resourceAwsWorkspacesWorkspaceRead(d, meta)
}

func resourceAwsWorkspacesWorkspaceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn
	ignoreTagsConfig := meta.(*AWSClient).IgnoreTagsConfig

	rawOutput, state, err := waiter.WorkspaceState(conn, d.Id())()
	if err != nil {
		return fmt.Errorf("error reading workspace (%s): %s", d.Id(), err)
	}
	if state == workspaces.WorkspaceStateTerminated {
		log.Printf("[WARN] workspace (%s) is not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	workspace := rawOutput.(*workspaces.Workspace)
	d.Set("bundle_id", workspace.BundleId)
	d.Set("directory_id", workspace.DirectoryId)
	d.Set("ip_address", workspace.IpAddress)
	d.Set("computer_name", workspace.ComputerName)
	d.Set("state", workspace.State)
	d.Set("root_volume_encryption_enabled", workspace.RootVolumeEncryptionEnabled)
	d.Set("user_name", workspace.UserName)
	d.Set("user_volume_encryption_enabled", workspace.UserVolumeEncryptionEnabled)
	d.Set("volume_encryption_key", workspace.VolumeEncryptionKey)
	if err := d.Set("workspace_properties", flattenWorkspaceProperties(workspace.WorkspaceProperties)); err != nil {
		return fmt.Errorf("error setting workspace properties: %s", err)
	}

	tags, err := keyvaluetags.WorkspacesListTags(conn, d.Id())
	if err != nil {
		return fmt.Errorf("error listing tags: %s", err)
	}

	if err := d.Set("tags", tags.IgnoreAws().IgnoreConfig(ignoreTagsConfig).Map()); err != nil {
		return fmt.Errorf("error setting tags: %s", err)
	}

	return nil
}

func resourceAwsWorkspacesWorkspaceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn

	// IMPORTANT: Only one workspace property could be changed in a time.
	// I've create AWS Support feature request to allow multiple properties modification in a time.
	// https://docs.aws.amazon.com/workspaces/latest/adminguide/modify-workspaces.html

	if d.HasChange("workspace_properties.0.compute_type_name") {
		if err := workspacePropertyUpdate("compute_type_name", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.root_volume_size_gib") {
		if err := workspacePropertyUpdate("root_volume_size_gib", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.running_mode") {
		if err := workspacePropertyUpdate("running_mode", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.running_mode_auto_stop_timeout_in_minutes") {
		if err := workspacePropertyUpdate("running_mode_auto_stop_timeout_in_minutes", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.user_volume_size_gib") {
		if err := workspacePropertyUpdate("user_volume_size_gib", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("tags") {
		o, n := d.GetChange("tags")
		if err := keyvaluetags.WorkspacesUpdateTags(conn, d.Id(), o, n); err != nil {
			return fmt.Errorf("error updating tags: %s", err)
		}
	}

	return resourceAwsWorkspacesWorkspaceRead(d, meta)
}

func resourceAwsWorkspacesWorkspaceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn

	err := workspaceDelete(d.Id(), conn)
	if err != nil {
		return err
	}

	return nil
}

func workspaceDelete(id string, conn *workspaces.WorkSpaces) error {
	log.Printf("[DEBUG] Terminating workspace %q", id)
	_, err := conn.TerminateWorkspaces(&workspaces.TerminateWorkspacesInput{
		TerminateWorkspaceRequests: []*workspaces.TerminateRequest{
			{
				WorkspaceId: aws.String(id),
			},
		},
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Waiting for workspace %q to be terminated", id)
	_, err = waiter.WorkspaceTerminated(conn, id)
	if err != nil {
		return fmt.Errorf("workspace %q was not terminated: %s", id, err)
	}
	log.Printf("[DEBUG] Workspace %q is terminated", id)

	return nil
}

func workspacePropertyUpdate(p string, conn *workspaces.WorkSpaces, d *schema.ResourceData) error {
	id := d.Id()

	log.Printf("[DEBUG] Modifying workspace %q %s property...", id, p)

	var wsp *workspaces.WorkspaceProperties

	switch p {
	case "compute_type_name":
		wsp = &workspaces.WorkspaceProperties{
			ComputeTypeName: aws.String(d.Get("workspace_properties.0.compute_type_name").(string)),
		}
	case "root_volume_size_gib":
		wsp = &workspaces.WorkspaceProperties{
			RootVolumeSizeGib: aws.Int64(int64(d.Get("workspace_properties.0.root_volume_size_gib").(int))),
		}
	case "running_mode":
		wsp = &workspaces.WorkspaceProperties{
			RunningMode: aws.String(d.Get("workspace_properties.0.running_mode").(string)),
		}
	case "running_mode_auto_stop_timeout_in_minutes":
		if d.Get("workspace_properties.0.running_mode") != workspaces.RunningModeAutoStop {
			log.Printf("[DEBUG] Property running_mode_auto_stop_timeout_in_minutes makes sense only for AUTO_STOP running mode")
			return nil
		}

		wsp = &workspaces.WorkspaceProperties{
			RunningModeAutoStopTimeoutInMinutes: aws.Int64(int64(d.Get("workspace_properties.0.running_mode_auto_stop_timeout_in_minutes").(int))),
		}
	case "user_volume_size_gib":
		wsp = &workspaces.WorkspaceProperties{
			UserVolumeSizeGib: aws.Int64(int64(d.Get("workspace_properties.0.user_volume_size_gib").(int))),
		}
	}

	_, err := conn.ModifyWorkspaceProperties(&workspaces.ModifyWorkspacePropertiesInput{
		WorkspaceId:         aws.String(id),
		WorkspaceProperties: wsp,
	})
	if err != nil {
		return fmt.Errorf("workspace %q %s property was not modified: %w", d.Id(), p, err)
	}

	log.Printf("[DEBUG] Waiting for workspace %q %s property to be modified...", d.Id(), p)
	_, err = waiter.WorkspaceUpdated(conn, d.Id())
	if err != nil {
		return fmt.Errorf("error modifying workspace %q property %q was not modified: %w", d.Id(), p, err)
	}
	log.Printf("[DEBUG] Workspace %q %s property is modified", d.Id(), p)

	return nil
}

func expandWorkspaceProperties(properties []interface{}) *workspaces.WorkspaceProperties {
	log.Printf("[DEBUG] Expand Workspace properties: %+v ", properties)

	if len(properties) == 0 || properties[0] == nil {
		return nil
	}

	p := properties[0].(map[string]interface{})

	return &workspaces.WorkspaceProperties{
		ComputeTypeName:                     aws.String(p["compute_type_name"].(string)),
		RootVolumeSizeGib:                   aws.Int64(int64(p["root_volume_size_gib"].(int))),
		RunningMode:                         aws.String(p["running_mode"].(string)),
		RunningModeAutoStopTimeoutInMinutes: aws.Int64(int64(p["running_mode_auto_stop_timeout_in_minutes"].(int))),
		UserVolumeSizeGib:                   aws.Int64(int64(p["user_volume_size_gib"].(int))),
	}
}

func flattenWorkspaceProperties(properties *workspaces.WorkspaceProperties) []map[string]interface{} {
	log.Printf("[DEBUG] Flatten workspace properties: %+v ", properties)

	if properties == nil {
		return []map[string]interface{}{}
	}

	return []map[string]interface{}{
		{
			"compute_type_name":                         aws.StringValue(properties.ComputeTypeName),
			"root_volume_size_gib":                      int(aws.Int64Value(properties.RootVolumeSizeGib)),
			"running_mode":                              aws.StringValue(properties.RunningMode),
			"running_mode_auto_stop_timeout_in_minutes": int(aws.Int64Value(properties.RunningModeAutoStopTimeoutInMinutes)),
			"user_volume_size_gib":                      int(aws.Int64Value(properties.UserVolumeSizeGib)),
		},
	}
}
<commit_msg>Update resource_aws_workspaces_workspace.go<commit_after>package aws

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/workspaces"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"

	"github.com/terraform-providers/terraform-provider-aws/aws/internal/keyvaluetags"
	"github.com/terraform-providers/terraform-provider-aws/aws/internal/service/workspaces/waiter"
)

func resourceAwsWorkspacesWorkspace() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsWorkspacesWorkspaceCreate,
		Read:   resourceAwsWorkspacesWorkspaceRead,
		Update: resourceAwsWorkspacesWorkspaceUpdate,
		Delete: resourceAwsWorkspacesWorkspaceDelete,
		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},
		Schema: map[string]*schema.Schema{
			"bundle_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"directory_id": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"ip_address": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"computer_name": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"state": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"root_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
				Default:  false,
			},
			"user_name": {
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"user_volume_encryption_enabled": {
				Type:     schema.TypeBool,
				Optional: true,
				ForceNew: true,
				Default:  false,
			},
			"volume_encryption_key": {
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},
			"workspace_properties": {
				Type:     schema.TypeList,
				MaxItems: 1,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"compute_type_name": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  workspaces.ComputeValue,
							ValidateFunc: validation.StringInSlice([]string{
								workspaces.ComputeValue,
								workspaces.ComputeStandard,
								workspaces.ComputePerformance,
								workspaces.ComputePower,
								workspaces.ComputePowerpro,
								workspaces.ComputeGraphics,
								workspaces.ComputeGraphicspro,
							}, false),
						},
						"root_volume_size_gib": {
							Type:     schema.TypeInt,
							Optional: true,
							Default:  80,
							ValidateFunc: validation.Any(
								validation.IntInSlice([]int{80}),
								validation.IntBetween(175, 2000),
							),
						},
						"running_mode": {
							Type:     schema.TypeString,
							Optional: true,
							Default:  workspaces.RunningModeAlwaysOn,
							ValidateFunc: validation.StringInSlice([]string{
								workspaces.RunningModeAlwaysOn,
								workspaces.RunningModeAutoStop,
							}, false),
						},
						"running_mode_auto_stop_timeout_in_minutes": {
							Type:     schema.TypeInt,
							Optional: true,
							Computed: true,
							ValidateFunc: func(v interface{}, k string) (ws []string, errors []error) {
								val := v.(int)
								if val%60 != 0 {
									errors = append(errors, fmt.Errorf(
										"%q should be configured in 60-minute intervals, got: %d", k, val))
								}
								return
							},
						},
						"user_volume_size_gib": {
							Type:     schema.TypeInt,
							Optional: true,
							Default:  10,
							ValidateFunc: validation.Any(
								validation.IntInSlice([]int{10, 50}),
								validation.IntBetween(100, 2000),
							),
						},
					},
				},
			},
			"tags": tagsSchema(),
		},
	}
}

func resourceAwsWorkspacesWorkspaceCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn

	tags := keyvaluetags.New(d.Get("tags").(map[string]interface{})).IgnoreAws().WorkspacesTags()

	input := &workspaces.WorkspaceRequest{
		BundleId:                    aws.String(d.Get("bundle_id").(string)),
		DirectoryId:                 aws.String(d.Get("directory_id").(string)),
		UserName:                    aws.String(d.Get("user_name").(string)),
		RootVolumeEncryptionEnabled: aws.Bool(d.Get("root_volume_encryption_enabled").(bool)),
		UserVolumeEncryptionEnabled: aws.Bool(d.Get("user_volume_encryption_enabled").(bool)),
		Tags:                        tags,
	}

	if v, ok := d.GetOk("volume_encryption_key"); ok {
		input.VolumeEncryptionKey = aws.String(v.(string))
	}

	input.WorkspaceProperties = expandWorkspaceProperties(d.Get("workspace_properties").([]interface{}))

	log.Printf("[DEBUG] Creating workspace...\n%#v\n", *input)
	resp, err := conn.CreateWorkspaces(&workspaces.CreateWorkspacesInput{
		Workspaces: []*workspaces.WorkspaceRequest{input},
	})
	if err != nil {
		return err
	}

	wsFail := resp.FailedRequests
	if len(wsFail) > 0 {
		return fmt.Errorf("workspace creation failed: %s", *wsFail[0].ErrorMessage)
	}

	workspaceID := aws.StringValue(resp.PendingRequests[0].WorkspaceId)

	log.Printf("[DEBUG] Waiting for workspace %q to be available...", workspaceID)
	_, err = waiter.WorkspaceAvailable(conn, workspaceID)
	if err != nil {
		return fmt.Errorf("workspace %q is not available: %s", workspaceID, err)
	}

	d.SetId(workspaceID)
	log.Printf("[DEBUG] Workspace %q is available", workspaceID)

	return resourceAwsWorkspacesWorkspaceRead(d, meta)
}

func resourceAwsWorkspacesWorkspaceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn
	ignoreTagsConfig := meta.(*AWSClient).IgnoreTagsConfig

	rawOutput, state, err := waiter.WorkspaceState(conn, d.Id())()
	if err != nil {
		return fmt.Errorf("error reading workspace (%s): %s", d.Id(), err)
	}
	if state == workspaces.WorkspaceStateTerminated {
		log.Printf("[WARN] workspace (%s) is not found, removing from state", d.Id())
		d.SetId("")
		return nil
	}

	workspace := rawOutput.(*workspaces.Workspace)
	d.Set("bundle_id", workspace.BundleId)
	d.Set("directory_id", workspace.DirectoryId)
	d.Set("ip_address", workspace.IpAddress)
	d.Set("computer_name", workspace.ComputerName)
	d.Set("state", workspace.State)
	d.Set("root_volume_encryption_enabled", workspace.RootVolumeEncryptionEnabled)
	d.Set("user_name", workspace.UserName)
	d.Set("user_volume_encryption_enabled", workspace.UserVolumeEncryptionEnabled)
	d.Set("volume_encryption_key", workspace.VolumeEncryptionKey)
	if err := d.Set("workspace_properties", flattenWorkspaceProperties(workspace.WorkspaceProperties)); err != nil {
		return fmt.Errorf("error setting workspace properties: %s", err)
	}

	tags, err := keyvaluetags.WorkspacesListTags(conn, d.Id())
	if err != nil {
		return fmt.Errorf("error listing tags: %s", err)
	}

	if err := d.Set("tags", tags.IgnoreAws().IgnoreConfig(ignoreTagsConfig).Map()); err != nil {
		return fmt.Errorf("error setting tags: %s", err)
	}

	return nil
}

func resourceAwsWorkspacesWorkspaceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn

	// IMPORTANT: Only one workspace property could be changed in a time.
	// I've create AWS Support feature request to allow multiple properties modification in a time.
	// https://docs.aws.amazon.com/workspaces/latest/adminguide/modify-workspaces.html

	if d.HasChange("workspace_properties.0.compute_type_name") {
		if err := workspacePropertyUpdate("compute_type_name", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.root_volume_size_gib") {
		if err := workspacePropertyUpdate("root_volume_size_gib", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.running_mode") {
		if err := workspacePropertyUpdate("running_mode", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.running_mode_auto_stop_timeout_in_minutes") || (d.Get("workspace_properties.0.running_mode") != "ALWAYS_ON") {
		if err := workspacePropertyUpdate("running_mode_auto_stop_timeout_in_minutes", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("workspace_properties.0.user_volume_size_gib") {
		if err := workspacePropertyUpdate("user_volume_size_gib", conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("tags") {
		o, n := d.GetChange("tags")
		if err := keyvaluetags.WorkspacesUpdateTags(conn, d.Id(), o, n); err != nil {
			return fmt.Errorf("error updating tags: %s", err)
		}
	}

	return resourceAwsWorkspacesWorkspaceRead(d, meta)
}

func resourceAwsWorkspacesWorkspaceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).workspacesconn

	err := workspaceDelete(d.Id(), conn)
	if err != nil {
		return err
	}

	return nil
}

func workspaceDelete(id string, conn *workspaces.WorkSpaces) error {
	log.Printf("[DEBUG] Terminating workspace %q", id)
	_, err := conn.TerminateWorkspaces(&workspaces.TerminateWorkspacesInput{
		TerminateWorkspaceRequests: []*workspaces.TerminateRequest{
			{
				WorkspaceId: aws.String(id),
			},
		},
	})
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] Waiting for workspace %q to be terminated", id)
	_, err = waiter.WorkspaceTerminated(conn, id)
	if err != nil {
		return fmt.Errorf("workspace %q was not terminated: %s", id, err)
	}
	log.Printf("[DEBUG] Workspace %q is terminated", id)

	return nil
}

func workspacePropertyUpdate(p string, conn *workspaces.WorkSpaces, d *schema.ResourceData) error {
	id := d.Id()

	log.Printf("[DEBUG] Modifying workspace %q %s property...", id, p)

	var wsp *workspaces.WorkspaceProperties

	switch p {
	case "compute_type_name":
		wsp = &workspaces.WorkspaceProperties{
			ComputeTypeName: aws.String(d.Get("workspace_properties.0.compute_type_name").(string)),
		}
	case "root_volume_size_gib":
		wsp = &workspaces.WorkspaceProperties{
			RootVolumeSizeGib: aws.Int64(int64(d.Get("workspace_properties.0.root_volume_size_gib").(int))),
		}
	case "running_mode":
		wsp = &workspaces.WorkspaceProperties{
			RunningMode: aws.String(d.Get("workspace_properties.0.running_mode").(string)),
		}
	case "running_mode_auto_stop_timeout_in_minutes":
		if d.Get("workspace_properties.0.running_mode") != workspaces.RunningModeAutoStop {
			log.Printf("[DEBUG] Property running_mode_auto_stop_timeout_in_minutes makes sense only for AUTO_STOP running mode")
			return nil
		}

		wsp = &workspaces.WorkspaceProperties{
			RunningModeAutoStopTimeoutInMinutes: aws.Int64(int64(d.Get("workspace_properties.0.running_mode_auto_stop_timeout_in_minutes").(int))),
		}
	case "user_volume_size_gib":
		wsp = &workspaces.WorkspaceProperties{
			UserVolumeSizeGib: aws.Int64(int64(d.Get("workspace_properties.0.user_volume_size_gib").(int))),
		}
	}

	_, err := conn.ModifyWorkspaceProperties(&workspaces.ModifyWorkspacePropertiesInput{
		WorkspaceId:         aws.String(id),
		WorkspaceProperties: wsp,
	})
	if err != nil {
		return fmt.Errorf("workspace %q %s property was not modified: %w", d.Id(), p, err)
	}

	log.Printf("[DEBUG] Waiting for workspace %q %s property to be modified...", d.Id(), p)
	_, err = waiter.WorkspaceUpdated(conn, d.Id())
	if err != nil {
		return fmt.Errorf("error modifying workspace %q property %q was not modified: %w", d.Id(), p, err)
	}
	log.Printf("[DEBUG] Workspace %q %s property is modified", d.Id(), p)

	return nil
}

func expandWorkspaceProperties(properties []interface{}) *workspaces.WorkspaceProperties {
	log.Printf("[DEBUG] Expand Workspace properties: %+v ", properties)

	if len(properties) == 0 || properties[0] == nil {
		return nil
	}

	p := properties[0].(map[string]interface{})

	if p["running_mode"] == "ALWAYS_ON" {
		return &workspaces.WorkspaceProperties{
			ComputeTypeName:   aws.String(p["compute_type_name"].(string)),
			RootVolumeSizeGib: aws.Int64(int64(p["root_volume_size_gib"].(int))),
			RunningMode:       aws.String(p["running_mode"].(string)),
			UserVolumeSizeGib: aws.Int64(int64(p["user_volume_size_gib"].(int))),
		}
	} else {
		return &workspaces.WorkspaceProperties{
			ComputeTypeName:                     aws.String(p["compute_type_name"].(string)),
			RootVolumeSizeGib:                   aws.Int64(int64(p["root_volume_size_gib"].(int))),
			RunningMode:                         aws.String(p["running_mode"].(string)),
			RunningModeAutoStopTimeoutInMinutes: aws.Int64(int64(p["running_mode_auto_stop_timeout_in_minutes"].(int))),
			UserVolumeSizeGib:                   aws.Int64(int64(p["user_volume_size_gib"].(int))),
		}
	}
}

func flattenWorkspaceProperties(properties *workspaces.WorkspaceProperties) []map[string]interface{} {
	log.Printf("[DEBUG] Flatten workspace properties: %+v ", properties)

	if properties == nil {
		return []map[string]interface{}{}
	}

	return []map[string]interface{}{
		{
			"compute_type_name":                         aws.StringValue(properties.ComputeTypeName),
			"root_volume_size_gib":                      int(aws.Int64Value(properties.RootVolumeSizeGib)),
			"running_mode":                              aws.StringValue(properties.RunningMode),
			"running_mode_auto_stop_timeout_in_minutes": int(aws.Int64Value(properties.RunningModeAutoStopTimeoutInMinutes)),
			"user_volume_size_gib":                      int(aws.Int64Value(properties.UserVolumeSizeGib)),
		},
	}
}
<|endoftext|>