go语言实现dcv端口转发

环境及目标

• 环境
  所有转发机器已经实现了端口和ip的转发。端口为默认端口。
  参考https://blog.csdn.net/wz122330/article/details/91980227

代码实现

package app

import (
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync/atomic"

	"com.wz/test/def"
	"com.wz/test/def"

	"gopkg.in/ini.v1"
)

type confParser struct {
	reader *ini.File // config reader
}

const (
	unavailStatus = "unavail"
	binCmd        = "/bin/bash"
	scriptName    = "get_host_status.sh"
	lineBreak     = "\n"
)

var (
	// Forwarded host file.
	forwardedHostFile = filepath.Join(sysconf.VarDir(), "forwarded.host")
	// For counter.
	count int32
)

// Get VNC console file.
func (m *Manager) GetVNCInfo(req *def.VNCReq) (string, error) {
	cwd := req.Cwd
	vncConsolePath := ""
	if filepath.IsAbs(cwd) {
		vncConsolePath = filepath.Join(cwd, def.Spooler, def.OperatorConsole)
	}
	if _, err := os.Stat(vncConsolePath); err != nil {
		log.Errorf("The job '%s' vnc file is not created, because of %v", req.JobID, err)
		return "", err
	}
	if m.config.DCV.Forward {
		if err := m.convertHostIP(vncConsolePath); err != nil {
			log.Errorf("Failed to convert host ip, vncConsolePath %s", vncConsolePath)
			return vncConsolePath, err
		}
	}
	return vncConsolePath, nil
}

func (m *Manager) convertHostIP(dcvFile string) error {
	section := def.Connect
	key := def.ForwardHost
	forwardHostsFile := m.config.DCV.HostsFile

	parser := confParser{}
	parser.load(dcvFile)

	lastHost, err := loadForwardHostFromFile(forwardedHostFile)
	if err != nil {
		log.Warnf("Failed to load forward host from file %s, because of %v", forwardedHostFile, err)
	}
	hosts, err := m.loadHostsToSlice(forwardHostsFile)
	if err != nil {
		log.Errorf("Failed to load hosts from file %s to slice, because of %v", forwardHostsFile, err)
		return err
	}
	h, err := m.chooseAvailHost(lastHost, hosts)
	if err != nil {
		log.Errorf("Failed to choose a available host, lastHost %s hosts %v, because of %v",
			lastHost, hosts, err)
		return err
	}
	if h != "" {
		log.Debugf("Choose a available host %s", h)
		if err := parser.setValue(dcvFile, section, key, h); err != nil {
			return err
		}
		if err := saveForwardHostToFile(forwardedHostFile, h); err != nil {
			log.Warnf("Failed to save host %s to file %s, because of %s", h, forwardedHostFile, err)
			return nil
		}
	} else {
		log.Errorf("All hosts in hosts file %s are unavailable, please check hosts's status",
			forwardHostsFile)
	}

	return nil
}

func getNextHost(lastHost string, hosts []string) string {
	atomic.AddInt32(&count, 1)
	log.Debugf("The counter is %d", count)
	idx := 0
	if lastHost == "" {
		return hosts[0]
	}
	for i, h := range hosts {
		if lastHost == h {
			idx = i + 1
			break
		} else {
			continue
		}
	}
	if idx < len(hosts) {
		return hosts[idx]
	} else {
		return hosts[0]
	}
}

func (m *Manager) chooseAvailHost(host string, hosts []string) (string, error) {
	if int(count) >= len(hosts) {
		return "", nil
	}
	log.Debugf("Last host is %s, all hosts are %v", host, hosts)
	nextHost := getNextHost(host, hosts)
	log.Debugf("NextHost is %s ", nextHost)
	available, err := m.checkHostAvailable(nextHost)
	if err != nil {
		return "", err
	} else {
		if available {
			return nextHost, nil
		} else {
			return m.chooseAvailHost(nextHost, hosts)
		}
	}
}

func (m *Manager) loadHostsToSlice(hostsFile string) ([]string, error) {
	cont, err := ioutil.ReadFile(hostsFile)
	if err != nil {
		return []string{}, err
	}
	// Remove the last one '\n'.
	str := strings.TrimSpace(string(cont))
	return strings.Split(str, lineBreak), nil
}

// Load conf file.
func (c *confParser) load(vncConsolePath string) {
	conf, err := ini.Load(vncConsolePath)
	if err != nil {
		log.Errorf("Failed to load vnc file %s, because of %v", vncConsolePath, err)
	}
	c.reader = conf
}

// Set value for conf file.
func (c *confParser) setValue(file, section, key, value string) error {
	c.reader.Section(section).Key(key).SetValue(value)
	return c.reader.SaveTo(file)
}

func (m *Manager) checkHostAvailable(ip string) (bool, error) {
	shellPath := filepath.Join(m.config.BaseDir, def.AppScriptDir, scriptName)
	log.Debugf("shellPath %s  %s", shellPath, ip)
	out, err := exec.Command(binCmd, shellPath, ip).CombinedOutput()
	if err != nil {
		return false, err
	}
	status := strings.TrimSpace(string(out))
	log.Debugf("Host is %s, it's status is %s", ip, status)
	if status == unavailStatus {
		return false, err
	}
	count = 0
	return true, nil
}

func loadForwardHostFromFile(file string) (string, error) {
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func saveForwardHostToFile(file, host string) error {
	log.Debugf("Save forwarding host %s to %s", host, file)
	if err := os.MkdirAll(filepath.Dir(file), def.DirPerm); err != nil {
		return err
	}
	return ioutil.WriteFile(file, []byte(host), def.FilePerm)
}