mirror of
https://github.com/ssut/payload-dumper-go.git
synced 2024-05-13 19:19:11 +00:00
feat: goroutine
This commit is contained in:
Suhun Han
parent
16b7bebd02
commit
228a048cbb
1
go.mod
1
go.mod
|
|
@ -6,5 +6,6 @@ require (
|
|||
github.com/dustin/go-humanize v1.0.0
|
||||
github.com/golang/protobuf v1.4.2
|
||||
github.com/spencercw/go-xz v0.0.0-20181128201811-c82a2123b492
|
||||
github.com/vbauerster/mpb/v5 v5.3.0
|
||||
google.golang.org/protobuf v1.25.0
|
||||
)
|
||||
|
|
|
|||
11
go.sum
11
go.sum
|
|
@ -1,5 +1,9 @@
|
|||
cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
|
||||
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
|
||||
github.com/VividCortex/ewma v1.1.1 h1:MnEK4VOv6n0RSY4vtRe3h11qjxL3+t0B8yOL8iMXdcM=
|
||||
github.com/VividCortex/ewma v1.1.1/go.mod h1:2Tkkvm3sRDVXaiyucHiACn4cqf7DpdyLvmxzcbUokwA=
|
||||
github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d h1:licZJFw2RwpHMqeKTCYkitsPqHNxTmd4SNR5r94FGM8=
|
||||
github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d/go.mod h1:asat636LX7Bqt5lYEZ27JNDcqxfjdBQuJ/MM4CN/Lzo=
|
||||
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
|
||||
github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
|
||||
github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
|
||||
|
|
@ -25,9 +29,14 @@ github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMyw
|
|||
github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
|
||||
github.com/google/go-cmp v0.5.0 h1:/QaMHBdZ26BB3SSst0Iwl10Epc+xhTquomWX0oZEB6w=
|
||||
github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
|
||||
github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0=
|
||||
github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
|
||||
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
|
||||
github.com/spencercw/go-xz v0.0.0-20181128201811-c82a2123b492 h1:8J9q7E8tGpVB84cBsMr+X160ECRqwYhkZ6KeaY9kN1I=
|
||||
github.com/spencercw/go-xz v0.0.0-20181128201811-c82a2123b492/go.mod h1:EvRrgz1GcjNV5yfN+ISxA4sxn255MimeGQ/ROJnQPtQ=
|
||||
github.com/vbauerster/mpb/v5 v5.3.0 h1:vgrEJjUzHaSZKDRRxul5Oh4C72Yy/5VEMb0em+9M0mQ=
|
||||
github.com/vbauerster/mpb/v5 v5.3.0/go.mod h1:4yTkvAb8Cm4eylAp6t0JRq6pXDkFJ4krUlDqWYkakAs=
|
||||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2 h1:VklqNMn3ovrHsnt90PveolxSbWFaJdECFbxSq0Mqo2M=
|
||||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
|
||||
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
|
||||
golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
|
||||
|
|
@ -43,6 +52,8 @@ golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJ
|
|||
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed h1:WBkVNH1zd9jg/dK4HCM4lNANnmd12EHC9z+LmcCG4ns=
|
||||
golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=
|
||||
|
|
|
|||
2
main.go
2
main.go
|
|
@ -8,6 +8,7 @@ import (
|
|||
"io/ioutil"
|
||||
"log"
|
||||
"os"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
|
@ -45,6 +46,7 @@ func extractPayloadBin(filename string) string {
|
|||
}
|
||||
|
||||
func main() {
|
||||
runtime.GOMAXPROCS(runtime.NumCPU())
|
||||
filename := os.Args[1]
|
||||
|
||||
if _, err := os.Stat(filename); os.IsNotExist(err) {
|
||||
|
|
|
|||
83
payload.go
83
payload.go
|
|
@ -9,13 +9,21 @@ import (
|
|||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
humanize "github.com/dustin/go-humanize"
|
||||
"github.com/golang/protobuf/proto"
|
||||
xz "github.com/spencercw/go-xz"
|
||||
"github.com/ssut/payload-dumper-go/chromeos_update_engine"
|
||||
"github.com/vbauerster/mpb/v5"
|
||||
"github.com/vbauerster/mpb/v5/decor"
|
||||
)
|
||||
|
||||
type request struct {
|
||||
partition *chromeos_update_engine.PartitionUpdate
|
||||
targetDirectory string
|
||||
}
|
||||
|
||||
// Payload is a new format for the Android OTA/Firmware update files since Android Oreo
|
||||
type Payload struct {
|
||||
Filename string
|
||||
|
|
@ -28,6 +36,10 @@ type Payload struct {
|
|||
metadataSize int64
|
||||
dataOffset int64
|
||||
initialized bool
|
||||
|
||||
requests chan *request
|
||||
workerWG sync.WaitGroup
|
||||
progress *mpb.Progress
|
||||
}
|
||||
|
||||
const payloadHeaderMagic = "CrAU"
|
||||
|
|
@ -192,10 +204,26 @@ func (p *Payload) readDataBlob(offset int64, length int64) ([]byte, error) {
|
|||
|
||||
func (p *Payload) Extract(partition *chromeos_update_engine.PartitionUpdate, out *os.File) error {
|
||||
name := partition.GetPartitionName()
|
||||
info := partition.GetNewPartitionInfo()
|
||||
totalOperations := len(partition.Operations)
|
||||
barName := fmt.Sprintf("%s (%s)", name, humanize.Bytes(info.GetSize()))
|
||||
bar := p.progress.AddBar(
|
||||
int64(totalOperations),
|
||||
mpb.PrependDecorators(
|
||||
decor.Name(barName, decor.WCSyncSpaceR),
|
||||
),
|
||||
mpb.AppendDecorators(
|
||||
decor.Percentage(),
|
||||
),
|
||||
)
|
||||
defer bar.SetTotal(0, true)
|
||||
|
||||
for _, operation := range partition.Operations {
|
||||
if len(operation.DstExtents) == 0 {
|
||||
return fmt.Errorf("Invalid operation.DstExtents for the partition %s", name)
|
||||
}
|
||||
bar.Increment()
|
||||
|
||||
e := operation.DstExtents[0]
|
||||
dataOffset := p.dataOffset + int64(operation.GetDataOffset())
|
||||
dataLength := int64(operation.GetDataLength())
|
||||
|
|
@ -210,7 +238,6 @@ func (p *Payload) Extract(partition *chromeos_update_engine.PartitionUpdate, out
|
|||
|
||||
switch operation.GetType() {
|
||||
case chromeos_update_engine.InstallOperation_REPLACE:
|
||||
fmt.Println("REPLACE")
|
||||
n, err := io.Copy(out, teeReader)
|
||||
if err != nil {
|
||||
return err
|
||||
|
|
@ -222,13 +249,11 @@ func (p *Payload) Extract(partition *chromeos_update_engine.PartitionUpdate, out
|
|||
break
|
||||
|
||||
case chromeos_update_engine.InstallOperation_REPLACE_XZ:
|
||||
fmt.Println("REPLACE_XZ")
|
||||
reader := xz.NewDecompressionReader(teeReader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Println(dataLength)
|
||||
_, err := io.Copy(out, &reader)
|
||||
if err != nil {
|
||||
return err
|
||||
|
|
@ -237,7 +262,6 @@ func (p *Payload) Extract(partition *chromeos_update_engine.PartitionUpdate, out
|
|||
break
|
||||
|
||||
case chromeos_update_engine.InstallOperation_REPLACE_BZ:
|
||||
fmt.Println("REPLACE_BZ")
|
||||
reader := bzip2.NewReader(teeReader)
|
||||
n, err := io.Copy(out, reader)
|
||||
if err != nil {
|
||||
|
|
@ -263,31 +287,48 @@ func (p *Payload) Extract(partition *chromeos_update_engine.PartitionUpdate, out
|
|||
return nil
|
||||
}
|
||||
|
||||
func (p *Payload) worker() {
|
||||
for req := range p.requests {
|
||||
partition := req.partition
|
||||
targetDirectory := req.targetDirectory
|
||||
|
||||
name := fmt.Sprintf("%s.img", partition.GetPartitionName())
|
||||
filepath := fmt.Sprintf("%s/%s", targetDirectory, name)
|
||||
file, err := os.OpenFile(filepath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0755)
|
||||
if err != nil {
|
||||
}
|
||||
if err := p.Extract(partition, file); err != nil {
|
||||
}
|
||||
|
||||
p.workerWG.Done()
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Payload) spawnExtractWorkers(n int) {
|
||||
for i := 0; i < n; i++ {
|
||||
go p.worker()
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Payload) ExtractAll(targetDirectory string) error {
|
||||
if !p.initialized {
|
||||
return errors.New("Payload has not been initialized")
|
||||
}
|
||||
p.progress = mpb.New()
|
||||
|
||||
total := len(p.deltaArchiveManifest.Partitions)
|
||||
for i, partition := range p.deltaArchiveManifest.Partitions {
|
||||
info := partition.GetNewPartitionInfo()
|
||||
name := fmt.Sprintf("%s.img", partition.GetPartitionName())
|
||||
p.requests = make(chan *request, 100)
|
||||
p.spawnExtractWorkers(4)
|
||||
|
||||
fmt.Printf("[%02d/%02d] Extracting: %s\n", i+1, total, fmt.Sprintf("%s (%s)", name, humanize.Bytes(*info.Size)))
|
||||
|
||||
filepath := fmt.Sprintf("%s/%s", targetDirectory, name)
|
||||
file, err := os.OpenFile(filepath, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0755)
|
||||
if err != nil {
|
||||
file.Close()
|
||||
return err
|
||||
for _, partition := range p.deltaArchiveManifest.Partitions {
|
||||
p.workerWG.Add(1)
|
||||
p.requests <- &request{
|
||||
partition: partition,
|
||||
targetDirectory: targetDirectory,
|
||||
}
|
||||
if err := p.Extract(partition, file); err != nil {
|
||||
file.Close()
|
||||
return err
|
||||
}
|
||||
file.Close()
|
||||
|
||||
}
|
||||
|
||||
p.workerWG.Wait()
|
||||
close(p.requests)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user