diff --git a/go.mod b/go.mod index b58fa2e..7a41f09 100644 --- a/go.mod +++ b/go.mod @@ -4,4 +4,4 @@ go 1.18 require github.com/logrusorgru/aurora v2.0.3+incompatible -require github.com/zhangpeihao/goamf v0.0.0-20140409082417-3ff2c19514a8 // indirect +require github.com/zhangpeihao/goamf v0.0.0-20140409082417-3ff2c19514a8 diff --git a/main.go b/main.go index 09edd60..c343219 100644 --- a/main.go +++ b/main.go @@ -26,47 +26,33 @@ type HDLConfig struct { var streamPathReg = regexp.MustCompile(`/(hdl/)?((.+)(\.flv)|(.+))`) func (config *HDLConfig) Update(override config.Config) { - override.Unmarshal(config) - if config.PullOnStart { - for streamPath, url := range config.AutoPullList { - if err := PullStream(streamPath, url); err != nil { - util.Println(err) - } - } - } if config.ListenAddr != "" || config.ListenAddrTLS != "" { - util.Print(Green("HDL Listen at "), BrightBlue(config.ListenAddr), BrightBlue(config.ListenAddrTLS)) + plugin.Infoln(Green("HDL Listen at "), BrightBlue(config.ListenAddr), BrightBlue(config.ListenAddrTLS)) config.Listen(plugin, config) } } -func (config *HDLConfig) API_pull(rw http.ResponseWriter, r *http.Request) { +func (config *HDLConfig) API_Pull(rw http.ResponseWriter, r *http.Request) { targetURL := r.URL.Query().Get("target") streamPath := r.URL.Query().Get("streamPath") - save := r.URL.Query().Get("save") - if err := PullStream(streamPath, targetURL); err == nil { - if save == "1" { + if config.PullStream(streamPath, Puller{RemoteURL: targetURL, Config: &config.Pull}) { + if r.URL.Query().Get("save") != "" { if config.AutoPullList == nil { config.AutoPullList = make(map[string]string) } config.AutoPullList[streamPath] = targetURL - if err = plugin.Save(); err != nil { - util.Println(err) + if err := plugin.Save(); err != nil { + plugin.Errorln(err) } } - rw.WriteHeader(200) - } else { - rw.WriteHeader(500) } } - -var hdlConfig = new(HDLConfig) -var plugin = InstallPlugin(hdlConfig) - -func init() { - plugin.HandleApi("/list", util.GetJsonHandler(FilterStreams[*HDLPuller], time.Second)) - plugin.HandleFunc("/", hdlConfig.ServeHTTP) +func (config *HDLConfig) API_List(rw http.ResponseWriter, r *http.Request) { + util.ReturnJson(FilterStreams[*HDLPuller], time.Second, rw, r) } +var Config = new(HDLConfig) +var plugin = InstallPlugin(Config) + func (config *HDLConfig) ServeHTTP(w http.ResponseWriter, r *http.Request) { parts := streamPathReg.FindStringSubmatch(r.RequestURI) if len(parts) == 0 { @@ -80,7 +66,7 @@ func (config *HDLConfig) ServeHTTP(w http.ResponseWriter, r *http.Request) { w.Header().Set("Transfer-Encoding", "chunked") w.Header().Set("Content-Type", "video/x-flv") sub := Subscriber{ID: r.RemoteAddr, Type: "FLV"} - if sub.Subscribe(stringPath, hdlConfig.Subscribe) { + if sub.Subscribe(stringPath, Config.Subscribe) { vt, at := sub.WaitVideoTrack(), sub.WaitAudioTrack() var buffer bytes.Buffer if _, err := amf.WriteString(&buffer, "onMetaData"); err != nil { diff --git a/pull.go b/pull.go index 4faf84c..dcf3ef5 100644 --- a/pull.go +++ b/pull.go @@ -4,7 +4,6 @@ import ( "bufio" "io" "net/http" - "net/url" "os" "strings" "time" @@ -67,61 +66,43 @@ type HDLPuller struct { at *track.UnknowAudio vt *track.UnknowVideo } +type FLVFile HDLPuller -func (puller *HDLPuller) OnStateChange(old StreamState, n StreamState) bool { - switch n { - case STATE_PUBLISHING: +func (puller *FLVFile) pull() { + (*HDLPuller)(puller).pull() +} + +func (puller *FLVFile) Pull(count int) { + if count == 0 { puller.at = puller.NewAudioTrack() puller.vt = puller.NewVideoTrack() - if puller.Type == "HDL Pull" { - if res, err := http.Get(puller.String()); err == nil { - puller.ReadCloser = res.Body - } else { - return false - } - } else { - if file, err := os.Open(puller.String()); err == nil { - puller.ReadCloser = file - } else { - file.Close() - return false - } - } - go puller.pull() - case STATE_WAITPUBLISH: - if hdlConfig.AutoReconnect { - if puller.Type == "HDL Pull" { - if res, err := http.Get(puller.String()); err == nil { - puller.ReadCloser = res.Body - } else { - return true - } - } else { - if file, err := os.Open(puller.String()); err == nil { - puller.ReadCloser = file - } else { - file.Close() - return true - } - go puller.pull() - } - } } - return true -} - -func PullStream(streamPath, address string) (err error) { - puller := &HDLPuller{} - puller.RemoteURL, err = url.Parse(address) - if err != nil { + if file, err := os.Open(puller.RemoteURL); err == nil { + puller.Reader = file + puller.Closer = file + } else { + file.Close() return } - if strings.HasPrefix(puller.Scheme, "http") { - puller.Type = "HDL Pull" - puller.Publish(streamPath, puller, hdlConfig.Publish) - } else { - puller.Type = "FLV File" - puller.Publish(streamPath, puller, hdlConfig.Publish) - } - return nil + puller.pull() +} + +func (puller *HDLPuller) Pull(count int) { + if count == 0 { + puller.at = puller.NewAudioTrack() + puller.vt = puller.NewVideoTrack() + } + if res, err := http.Get(puller.RemoteURL); err == nil { + puller.Reader = res.Body + puller.Closer = res.Body + } + puller.pull() +} + +func (config *HDLConfig) PullStream(streamPath string, puller Puller) bool { + if strings.HasPrefix(puller.RemoteURL, "http") { + return puller.Publish(streamPath, &HDLPuller{Puller: puller}, Config.Publish) + } else { + return puller.Publish(streamPath, &FLVFile{Puller: puller}, Config.Publish) + } }